andygrove commented on code in PR #2807:
URL:
https://github.com/apache/datafusion-ballista/pull/2807#discussion_r4519596854
##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -426,55 +505,25 @@ impl ExecutionPlan for ShuffleReaderExec {
// caller asked for.
let partition = if self.broadcast { 0 } else { partition };
- let config = context.session_config();
- let batch_size = config.batch_size();
-
- if config.ballista_shuffle_reader_force_remote_read() {
- debug!(
- "All shuffle partitions will be read as remote partitions! To
disable this behavior set: `{}=false`",
- crate::config::BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ
+ if self.broadcast
+ && let Some(cache) = &self.broadcast_cache
+ && context
+ .session_config()
+ .ballista_config()
+ .shuffle_cache_broadcast_reads()
+ && let Some(first) = self.partition[0].first()
+ {
+ let key = BroadcastKey::new(
+ first.partition_id.job_id.as_str(),
+ self.stage_id,
+ &self.partition[0],
+ );
+ let reader = self.clone();
+ return Ok(
+ cache.stream(key, self.schema.clone(), move || reader.fetch(0,
context))
);
}
-
- log::debug!(
- "ShuffleReaderExec::execute({task_id}) max_request_num: {},
max_message_size: {}",
- config.ballista_shuffle_reader_maximum_concurrent_requests(),
- config.ballista_grpc_client_max_message_size()
- );
- let mut partition_locations: Vec<PartitionLocation> =
- self.partition[partition].clone();
- // Shuffle partitions for evenly send fetching partition requests to
avoid hot executors within multiple tasks
- partition_locations.shuffle(&mut rng());
-
- let work_dir = self
- .work_dir
- .as_ref()
- .ok_or(DataFusionError::Configuration(
- "ShuffleReader work dir should have been set by
executor".to_owned(),
- ))?;
-
- let read_metrics = ShuffleReadMetrics::new(partition, &self.metrics);
-
- let response_receiver = send_fetch_partitions(
- work_dir,
- partition_locations,
- config,
- self.client_pool.clone(),
- read_metrics,
- );
-
- let input_stream = Box::pin(RecordBatchStreamAdapter::new(
- self.schema.clone(),
- response_receiver.try_flatten(),
- ));
-
- Ok(Box::pin(CoalescedShuffleReaderStream::new(
- input_stream,
- batch_size,
- None,
- &self.metrics,
- partition,
- )))
+ self.fetch(partition, context)
Review Comment:
Done, it's an if/else now.
##########
ballista/core/src/execution_plans/shuffle_reader.rs:
##########
@@ -426,55 +505,25 @@ impl ExecutionPlan for ShuffleReaderExec {
// caller asked for.
let partition = if self.broadcast { 0 } else { partition };
- let config = context.session_config();
- let batch_size = config.batch_size();
-
- if config.ballista_shuffle_reader_force_remote_read() {
- debug!(
- "All shuffle partitions will be read as remote partitions! To
disable this behavior set: `{}=false`",
- crate::config::BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ
+ if self.broadcast
+ && let Some(cache) = &self.broadcast_cache
+ && context
+ .session_config()
+ .ballista_config()
+ .shuffle_cache_broadcast_reads()
+ && let Some(first) = self.partition[0].first()
+ {
+ let key = BroadcastKey::new(
+ first.partition_id.job_id.as_str(),
+ self.stage_id,
+ &self.partition[0],
+ );
+ let reader = self.clone();
+ return Ok(
+ cache.stream(key, self.schema.clone(), move || reader.fetch(0,
context))
Review Comment:
Good catch, cache hits recorded zero `output_rows`. Tasks served from the
cache now record their own output rows, so every task's reader reports what it
produced. Fetch metrics stay at zero on a hit since nothing was fetched. Added
`broadcast_cache_hits_record_output_rows`.
##########
ballista/core/src/execution_plans/broadcast_cache.rs:
##########
@@ -0,0 +1,355 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Executor-wide cache of decoded broadcast shuffle reads.
+//!
+//! A broadcast hash join's build side is a broadcast [`ShuffleReaderExec`]
+//! that reads every output partition of the upstream stage. Each task of the
+//! joining stage runs its own copy of the plan, so without a cache every task
+//! on an executor fetches and decodes the same data: at 128 partitions on two
+//! executors, a broadcast stage of 16 tasks fetched 16,384 blocks and decoded
+//! the same build side 16 times. Spark broadcasts a relation once per
+//! executor.
+//!
+//! [`BroadcastReadCache`] keeps the decoded batches of each broadcast read so
+//! that the first task on an executor fetches them and the others reuse them.
+//! Tasks that ask while the first fetch is running wait for it. A failed fetch
+//! is not cached: every task that asks retries it and reports its own error,
+//! so fetch-failure recovery works as before.
+//!
+//! Entries are dropped when unused for [`BroadcastReadCache::idle_ttl`], when
+//! their job's data is removed, and are not kept at all above
+//! [`BroadcastReadCache::max_entry_bytes`]. The batches are not charged to a
+//! task's memory pool, as the fetched batches of a single read never were: the
+//! pool sees each task's hash table, which is built from them.
+//!
+//! [`ShuffleReaderExec`]: super::ShuffleReaderExec
+
+use std::collections::HashMap;
+use std::hash::{Hash, Hasher};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::common::utils::memory::RecordBatchMemoryCounter;
+use datafusion::error::Result;
+use datafusion::physical_plan::SendableRecordBatchStream;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use futures::{StreamExt, TryStreamExt};
+use log::debug;
+use tokio::sync::OnceCell;
+
+use crate::serde::scheduler::PartitionLocation;
+
+/// Identifies one broadcast read: the job, the upstream stage, and the exact
+/// locations read, so a re-run upstream stage (new locations) is a new entry.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct BroadcastKey {
+ job_id: String,
+ stage_id: usize,
+ locations: u64,
+}
+
+impl BroadcastKey {
+ /// Builds the key for reading `locations` of `stage_id` in `job_id`.
+ pub fn new(job_id: &str, stage_id: usize, locations: &[PartitionLocation])
-> Self {
+ let mut hasher = std::collections::hash_map::DefaultHasher::new();
+ locations.len().hash(&mut hasher);
+ for loc in locations {
+ loc.executor_meta.id.hash(&mut hasher);
+ loc.map_partition_id.hash(&mut hasher);
+ loc.partition_id.stage_id.hash(&mut hasher);
+ loc.partition_id.partition_id.hash(&mut hasher);
+ loc.file_id.hash(&mut hasher);
+ loc.is_sort_shuffle.hash(&mut hasher);
+ }
+ Self {
+ job_id: job_id.to_string(),
+ stage_id,
+ locations: hasher.finish(),
+ }
+ }
+}
+
+#[derive(Debug)]
+struct Entry {
+ batches: OnceCell<Arc<Vec<RecordBatch>>>,
+ last_used: Mutex<Instant>,
+}
+
+/// Executor-wide cache of decoded broadcast shuffle reads. See the
+/// [module documentation](self).
+#[derive(Debug)]
+pub struct BroadcastReadCache {
+ entries: Mutex<HashMap<BroadcastKey, Arc<Entry>>>,
+ idle_ttl: Duration,
+ max_entry_bytes: usize,
+}
+
+impl Default for BroadcastReadCache {
+ fn default() -> Self {
+ Self::new(Duration::from_secs(30), 1024 * 1024 * 1024)
+ }
+}
+
+impl BroadcastReadCache {
+ /// Creates a cache that drops entries unused for `idle_ttl` and does not
+ /// keep reads larger than `max_entry_bytes`.
+ pub fn new(idle_ttl: Duration, max_entry_bytes: usize) -> Self {
+ Self {
+ entries: Mutex::new(HashMap::new()),
+ idle_ttl,
+ max_entry_bytes,
+ }
+ }
+
+ /// How long an unused entry is kept.
+ pub fn idle_ttl(&self) -> Duration {
+ self.idle_ttl
+ }
+
+ /// The largest read, in bytes, that is kept for reuse.
+ pub fn max_entry_bytes(&self) -> usize {
+ self.max_entry_bytes
+ }
+
+ /// Number of reads currently held.
+ pub fn len(&self) -> usize {
+ self.entries.lock().unwrap().len()
+ }
+
+ /// Whether no reads are held.
+ pub fn is_empty(&self) -> bool {
+ self.len() == 0
+ }
+
+ /// Drops every read of `job_id`.
+ pub fn remove_job(&self, job_id: &str) {
+ self.entries
+ .lock()
+ .unwrap()
+ .retain(|key, _| key.job_id != job_id);
+ }
+
+ /// Returns a stream over the batches of the read `key`, fetching them with
+ /// `fetch` if no other task on this executor has, or waiting for the task
+ /// that is fetching them.
+ pub fn stream<F>(
+ self: &Arc<Self>,
+ key: BroadcastKey,
+ schema: SchemaRef,
+ fetch: F,
+ ) -> SendableRecordBatchStream
+ where
+ F: FnOnce() -> Result<SendableRecordBatchStream> + Send + 'static,
+ {
+ let entry = self.entry(&key);
+ let cache = Arc::clone(self);
+ let batches = async move {
+ let batches = entry
+ .batches
+ .get_or_try_init(|| async {
+ let batches: Vec<RecordBatch> =
fetch()?.try_collect().await?;
+ let mut counter = RecordBatchMemoryCounter::new();
+ let bytes: usize =
+ batches.iter().map(|b| counter.count_batch(b)).sum();
+ debug!(
+ "Fetched broadcast read of stage {} ({} batches,
{bytes} bytes)",
+ key.stage_id,
+ batches.len()
+ );
+ if bytes > cache.max_entry_bytes {
+ // Too large to hold between tasks: this task and any
+ // waiting for it still share it, later ones fetch
again.
+ cache.entries.lock().unwrap().remove(&key);
+ }
+ Ok::<_,
datafusion::error::DataFusionError>(Arc::new(batches))
+ })
+ .await?;
Review Comment:
Yes, it was stamped before the fetch, so a fetch slower than the idle TTL
left an entry that looked idle right away and got swept. It's now set in
`stream` once the batches are ready, which covers the fetching task and the
ones waiting on it. Added `a_slow_fetch_is_not_idle_when_it_finishes`.
##########
ballista/core/src/execution_plans/broadcast_cache.rs:
##########
@@ -0,0 +1,355 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Executor-wide cache of decoded broadcast shuffle reads.
+//!
+//! A broadcast hash join's build side is a broadcast [`ShuffleReaderExec`]
+//! that reads every output partition of the upstream stage. Each task of the
+//! joining stage runs its own copy of the plan, so without a cache every task
+//! on an executor fetches and decodes the same data: at 128 partitions on two
+//! executors, a broadcast stage of 16 tasks fetched 16,384 blocks and decoded
+//! the same build side 16 times. Spark broadcasts a relation once per
+//! executor.
+//!
+//! [`BroadcastReadCache`] keeps the decoded batches of each broadcast read so
+//! that the first task on an executor fetches them and the others reuse them.
+//! Tasks that ask while the first fetch is running wait for it. A failed fetch
+//! is not cached: every task that asks retries it and reports its own error,
+//! so fetch-failure recovery works as before.
+//!
+//! Entries are dropped when unused for [`BroadcastReadCache::idle_ttl`], when
+//! their job's data is removed, and are not kept at all above
+//! [`BroadcastReadCache::max_entry_bytes`]. The batches are not charged to a
+//! task's memory pool, as the fetched batches of a single read never were: the
+//! pool sees each task's hash table, which is built from them.
+//!
+//! [`ShuffleReaderExec`]: super::ShuffleReaderExec
+
+use std::collections::HashMap;
+use std::hash::{Hash, Hasher};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::common::utils::memory::RecordBatchMemoryCounter;
+use datafusion::error::Result;
+use datafusion::physical_plan::SendableRecordBatchStream;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use futures::{StreamExt, TryStreamExt};
+use log::debug;
+use tokio::sync::OnceCell;
+
+use crate::serde::scheduler::PartitionLocation;
+
+/// Identifies one broadcast read: the job, the upstream stage, and the exact
+/// locations read, so a re-run upstream stage (new locations) is a new entry.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct BroadcastKey {
+ job_id: String,
+ stage_id: usize,
+ locations: u64,
+}
+
+impl BroadcastKey {
+ /// Builds the key for reading `locations` of `stage_id` in `job_id`.
+ pub fn new(job_id: &str, stage_id: usize, locations: &[PartitionLocation])
-> Self {
+ let mut hasher = std::collections::hash_map::DefaultHasher::new();
+ locations.len().hash(&mut hasher);
+ for loc in locations {
+ loc.executor_meta.id.hash(&mut hasher);
+ loc.map_partition_id.hash(&mut hasher);
+ loc.partition_id.stage_id.hash(&mut hasher);
+ loc.partition_id.partition_id.hash(&mut hasher);
+ loc.file_id.hash(&mut hasher);
+ loc.is_sort_shuffle.hash(&mut hasher);
+ }
+ Self {
+ job_id: job_id.to_string(),
+ stage_id,
+ locations: hasher.finish(),
+ }
+ }
+}
+
+#[derive(Debug)]
+struct Entry {
+ batches: OnceCell<Arc<Vec<RecordBatch>>>,
+ last_used: Mutex<Instant>,
+}
+
+/// Executor-wide cache of decoded broadcast shuffle reads. See the
+/// [module documentation](self).
+#[derive(Debug)]
+pub struct BroadcastReadCache {
+ entries: Mutex<HashMap<BroadcastKey, Arc<Entry>>>,
+ idle_ttl: Duration,
+ max_entry_bytes: usize,
+}
+
+impl Default for BroadcastReadCache {
+ fn default() -> Self {
+ Self::new(Duration::from_secs(30), 1024 * 1024 * 1024)
+ }
+}
+
+impl BroadcastReadCache {
+ /// Creates a cache that drops entries unused for `idle_ttl` and does not
+ /// keep reads larger than `max_entry_bytes`.
+ pub fn new(idle_ttl: Duration, max_entry_bytes: usize) -> Self {
+ Self {
+ entries: Mutex::new(HashMap::new()),
+ idle_ttl,
+ max_entry_bytes,
+ }
+ }
+
+ /// How long an unused entry is kept.
+ pub fn idle_ttl(&self) -> Duration {
+ self.idle_ttl
+ }
+
+ /// The largest read, in bytes, that is kept for reuse.
+ pub fn max_entry_bytes(&self) -> usize {
+ self.max_entry_bytes
+ }
+
+ /// Number of reads currently held.
+ pub fn len(&self) -> usize {
+ self.entries.lock().unwrap().len()
+ }
+
+ /// Whether no reads are held.
+ pub fn is_empty(&self) -> bool {
+ self.len() == 0
+ }
+
+ /// Drops every read of `job_id`.
+ pub fn remove_job(&self, job_id: &str) {
+ self.entries
+ .lock()
+ .unwrap()
+ .retain(|key, _| key.job_id != job_id);
+ }
+
+ /// Returns a stream over the batches of the read `key`, fetching them with
+ /// `fetch` if no other task on this executor has, or waiting for the task
+ /// that is fetching them.
+ pub fn stream<F>(
+ self: &Arc<Self>,
+ key: BroadcastKey,
+ schema: SchemaRef,
+ fetch: F,
+ ) -> SendableRecordBatchStream
+ where
+ F: FnOnce() -> Result<SendableRecordBatchStream> + Send + 'static,
+ {
+ let entry = self.entry(&key);
+ let cache = Arc::clone(self);
+ let batches = async move {
+ let batches = entry
+ .batches
+ .get_or_try_init(|| async {
+ let batches: Vec<RecordBatch> =
fetch()?.try_collect().await?;
+ let mut counter = RecordBatchMemoryCounter::new();
+ let bytes: usize =
+ batches.iter().map(|b| counter.count_batch(b)).sum();
+ debug!(
+ "Fetched broadcast read of stage {} ({} batches,
{bytes} bytes)",
+ key.stage_id,
+ batches.len()
+ );
+ if bytes > cache.max_entry_bytes {
+ // Too large to hold between tasks: this task and any
+ // waiting for it still share it, later ones fetch
again.
+ cache.entries.lock().unwrap().remove(&key);
+ }
+ Ok::<_,
datafusion::error::DataFusionError>(Arc::new(batches))
+ })
+ .await?;
+ Ok::<_, datafusion::error::DataFusionError>(futures::stream::iter(
+ batches.iter().cloned().map(Ok).collect::<Vec<_>>(),
+ ))
+ };
+ Box::pin(RecordBatchStreamAdapter::new(
+ schema,
+ futures::stream::once(batches).try_flatten().boxed(),
+ ))
+ }
+
+ /// The entry for `key`, created if missing, after dropping idle entries.
+ fn entry(&self, key: &BroadcastKey) -> Arc<Entry> {
+ let now = Instant::now();
+ let mut entries = self.entries.lock().unwrap();
+ entries.retain(|k, e| {
+ k == key || now - *e.last_used.lock().unwrap() < self.idle_ttl
+ });
+ let entry = entries
+ .entry(key.clone())
+ .or_insert_with(|| {
+ Arc::new(Entry {
+ batches: OnceCell::new(),
+ last_used: Mutex::new(now),
+ })
+ })
+ .clone();
+ *entry.last_used.lock().unwrap() = now;
Review Comment:
Moved, see the reply 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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
