This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 59cd7528 feat: Feat force shuffle reader to read all files using arrow
flight client (#1310)
59cd7528 is described below
commit 59cd7528415d57145609f8b98f8e4d0bf722fd87
Author: Marko Milenković <[email protected]>
AuthorDate: Wed Sep 10 16:52:15 2025 +0100
feat: Feat force shuffle reader to read all files using arrow flight client
(#1310)
* add ability to force shuffle read to fetch all files
like they are remote.
new configuration option has been exposed:
`ballista.shuffle.force_remote_read=false`
* minor test setup issue
---
ballista/client/tests/common/mod.rs | 12 +++-
ballista/client/tests/context_checks.rs | 58 ++++++++++++++++
ballista/core/src/config.rs | 17 +++++
.../core/src/execution_plans/shuffle_reader.rs | 77 ++++++++++++++++++++--
ballista/core/src/extension.rs | 40 ++++++++++-
5 files changed, 194 insertions(+), 10 deletions(-)
diff --git a/ballista/client/tests/common/mod.rs
b/ballista/client/tests/common/mod.rs
index 800c21ea..85cb3fbc 100644
--- a/ballista/client/tests/common/mod.rs
+++ b/ballista/client/tests/common/mod.rs
@@ -236,13 +236,21 @@ pub async fn remote_context() -> SessionContext {
#[allow(dead_code)]
pub async fn standalone_context_with_state() -> SessionContext {
- let state = SessionStateBuilder::new().with_default_features().build();
+ let config = SessionConfig::new_with_ballista();
+ let state = SessionStateBuilder::new()
+ .with_config(config)
+ .with_default_features()
+ .build();
SessionContext::standalone_with_state(state).await.unwrap()
}
#[allow(dead_code)]
pub async fn remote_context_with_state() -> SessionContext {
- let state = SessionStateBuilder::new().with_default_features().build();
+ let config = SessionConfig::new_with_ballista();
+ let state = SessionStateBuilder::new()
+ .with_config(config)
+ .with_default_features()
+ .build();
let (host, port) = setup_test_cluster_with_state(state.clone()).await;
SessionContext::remote_with_state(&format!("df://{host}:{port}"), state)
.await
diff --git a/ballista/client/tests/context_checks.rs
b/ballista/client/tests/context_checks.rs
index a936f143..a4392f77 100644
--- a/ballista/client/tests/context_checks.rs
+++ b/ballista/client/tests/context_checks.rs
@@ -814,4 +814,62 @@ mod supported {
Ok(())
}
+
+ #[rstest]
+ #[case::standalone(standalone_context())]
+ #[case::remote(remote_context())]
+ #[tokio::test]
+ async fn should_force_local_read(
+ #[future(awt)]
+ #[case]
+ ctx: SessionContext,
+ test_data: String,
+ ) -> datafusion::error::Result<()> {
+ ctx.register_parquet(
+ "test",
+ &format!("{test_data}/alltypes_plain.parquet"),
+ Default::default(),
+ )
+ .await?;
+
+ ctx.sql("SET ballista.shuffle.force_remote_read = true")
+ .await?
+ .show()
+ .await?;
+
+ let result = ctx
+ .sql("select name, value from information_schema.df_settings where
name like 'ballista.shuffle.force_remote_read' order by name limit 1")
+ .await?
+ .collect()
+ .await?;
+
+ let expected = [
+ "+------------------------------------+-------+",
+ "| name | value |",
+ "+------------------------------------+-------+",
+ "| ballista.shuffle.force_remote_read | true |",
+ "+------------------------------------+-------+",
+ ];
+
+ assert_batches_eq!(expected, &result);
+
+ let expected = [
+ "+------------+----------+",
+ "| string_col | count(*) |",
+ "+------------+----------+",
+ "| 30 | 1 |",
+ "| 31 | 2 |",
+ "+------------+----------+",
+ ];
+
+ let result = ctx
+ .sql("select string_col, count(*) from test where id > 4 group by
string_col order by string_col")
+ .await?
+ .collect()
+ .await?;
+
+ assert_batches_eq!(expected, &result);
+
+ Ok(())
+ }
}
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index b5b0ea23..823d1d16 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -34,6 +34,8 @@ pub const BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE: &str =
"ballista.grpc_client_max_message_size";
pub const BALLISTA_SHUFFLE_READER_MAX_REQUESTS: &str =
"ballista.shuffle.max_concurrent_read_requests";
+pub const BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ: &str =
+ "ballista.shuffle.force_remote_read";
pub type ParseResult<T> = result::Result<T, String>;
use std::sync::LazyLock;
@@ -54,6 +56,11 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String,
ConfigEntry>> = LazyLock::new(||
"Maximum concurrent requests shuffle reader can
process".to_string(),
DataType::UInt64,
Some((64).to_string())),
+ ConfigEntry::new(BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ.to_string(),
+ "Forces the shuffle reader to always read partitions
via the Arrow Flight client, even when partitions are local to the
node.".to_string(),
+ DataType::Boolean,
+ Some((false).to_string())),
+
];
entries
.into_iter()
@@ -175,6 +182,16 @@ impl BallistaConfig {
self.get_usize_setting(BALLISTA_SHUFFLE_READER_MAX_REQUESTS)
}
+ /// Forces the shuffle reader to always read partitions via the Arrow
Flight client,
+ /// even when partitions are local to the node.
+ ///
+ /// Enabling forced remote read may significantly reduce performance,
+ /// as all partition reads will go through the Arrow Flight client even
for local data.
+ /// Use only when necessary, like in tests
+ pub fn shuffle_reader_force_remote_read(&self) -> bool {
+ self.get_bool_setting(BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ)
+ }
+
fn get_usize_setting(&self, key: &str) -> usize {
if let Some(v) = self.settings.get(key) {
// infallible because we validate all configs in the constructor
diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs
b/ballista/core/src/execution_plans/shuffle_reader.rs
index b0c0f04f..53624d17 100644
--- a/ballista/core/src/execution_plans/shuffle_reader.rs
+++ b/ballista/core/src/execution_plans/shuffle_reader.rs
@@ -160,6 +160,14 @@ impl ExecutionPlan for ShuffleReaderExec {
let max_request_num =
config.ballista_shuffle_reader_maximum_concurrent_requests();
let max_message_size = config.ballista_grpc_client_max_message_size();
+ let force_remote_read =
config.ballista_shuffle_reader_force_remote_read();
+
+ if 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
+ );
+ }
log::debug!(
"ShuffleReaderExec::execute({task_id}) max_request_num:
{max_request_num}, max_message_size: {max_message_size}"
@@ -180,8 +188,12 @@ impl ExecutionPlan for ShuffleReaderExec {
.collect();
// Shuffle partitions for evenly send fetching partition requests to
avoid hot executors within multiple tasks
partition_locations.shuffle(&mut rng());
- let response_receiver =
- send_fetch_partitions(partition_locations, max_request_num,
max_message_size);
+ let response_receiver = send_fetch_partitions(
+ partition_locations,
+ max_request_num,
+ max_message_size,
+ force_remote_read,
+ );
let result = RecordBatchStreamAdapter::new(
Arc::new(self.schema.as_ref().clone()),
@@ -352,18 +364,35 @@ impl Stream for AbortableReceiverStream {
.map_err(|e| ArrowError::ExternalError(Box::new(e)))
}
}
+/// Splits the provided partition locations into local and remote partitions.
+/// Local partitions are read directly from local Arrow IPC files,
+/// while remote partitions are fetched using the Arrow Flight client.
+/// If `force_remote_read` is true, all partitions are treated as remote.
+fn local_remote_read_split(
+ partition_locations: Vec<PartitionLocation>,
+ force_remote_read: bool,
+) -> (Vec<PartitionLocation>, Vec<PartitionLocation>) {
+ if !force_remote_read {
+ partition_locations
+ .into_iter()
+ .partition(check_is_local_location)
+ } else {
+ (vec![], partition_locations)
+ }
+}
fn send_fetch_partitions(
partition_locations: Vec<PartitionLocation>,
max_request_num: usize,
max_message_size: usize,
+ force_remote_read: bool,
) -> AbortableReceiverStream {
let (response_sender, response_receiver) = mpsc::channel(max_request_num);
let semaphore = Arc::new(Semaphore::new(max_request_num));
let mut spawned_tasks: Vec<SpawnedTask<()>> = vec![];
- let (local_locations, remote_locations): (Vec<_>, Vec<_>) =
partition_locations
- .into_iter()
- .partition(check_is_local_location);
+
+ let (local_locations, remote_locations): (Vec<_>, Vec<_>) =
+ local_remote_read_split(partition_locations, force_remote_read);
debug!(
"local shuffle file counts:{}, remote shuffle file count:{}.",
@@ -854,6 +883,36 @@ mod tests {
}
}
+ // tests if force remote read configuration option will
+ // qualify all partitions as remote
+ #[tokio::test]
+ async fn test_remote_local_read() {
+ let schema = get_test_partition_schema();
+ let data_array = Int32Array::from(vec![1]);
+ let batch =
+ RecordBatch::try_new(Arc::new(schema.clone()),
vec![Arc::new(data_array)])
+ .unwrap();
+ let tmp_dir = tempdir().unwrap();
+ let file_path = tmp_dir.path().join("shuffle_data");
+ let file = File::create(&file_path).unwrap();
+ let mut writer = StreamWriter::try_new(file, &schema).unwrap();
+ writer.write(&batch).unwrap();
+ writer.finish().unwrap();
+
+ let partition_locations =
+ get_test_partition_locations(1,
file_path.to_str().unwrap().to_string());
+
+ let (local, remote) =
local_remote_read_split(partition_locations.clone(), false);
+
+ assert!(!local.is_empty());
+ assert!(remote.is_empty());
+
+ let (local, remote) = local_remote_read_split(partition_locations,
true);
+
+ assert!(local.is_empty());
+ assert!(!remote.is_empty());
+ }
+
async fn test_send_fetch_partitions(max_request_num: usize, partition_num:
usize) {
let schema = get_test_partition_schema();
let data_array = Int32Array::from(vec![1]);
@@ -872,8 +931,12 @@ mod tests {
file_path.to_str().unwrap().to_string(),
);
- let response_receiver =
- send_fetch_partitions(partition_locations, max_request_num, 4 *
1024 * 1024);
+ let response_receiver = send_fetch_partitions(
+ partition_locations,
+ max_request_num,
+ 4 * 1024 * 1024,
+ false,
+ );
let stream = RecordBatchStreamAdapter::new(
Arc::new(schema),
diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs
index f1c0c200..5c4ba944 100644
--- a/ballista/core/src/extension.rs
+++ b/ballista/core/src/extension.rs
@@ -17,7 +17,8 @@
use crate::config::{
BallistaConfig, BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE, BALLISTA_JOB_NAME,
- BALLISTA_SHUFFLE_READER_MAX_REQUESTS, BALLISTA_STANDALONE_PARALLELISM,
+ BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ,
BALLISTA_SHUFFLE_READER_MAX_REQUESTS,
+ BALLISTA_STANDALONE_PARALLELISM,
};
use crate::planner::BallistaQueryPlanner;
use crate::serde::protobuf::KeyValuePair;
@@ -95,6 +96,21 @@ pub trait SessionConfigExt {
/// Returns parallelism of standalone cluster
fn ballista_standalone_parallelism(&self) -> usize;
+
+ /// Forces the shuffle reader to always read partitions via the Arrow
Flight client,
+ /// even when partitions are local to the node.
+ fn ballista_shuffle_reader_force_remote_read(&self) -> bool;
+ /// Forces the shuffle reader to always read partitions via the Arrow
Flight client,
+ /// even when partitions are local to the node.
+ ///
+ /// Enabling forced remote read may significantly reduce performance,
+ /// as all partition reads will go through the Arrow Flight client even
for local data.
+ /// Use only when necessary, like in tests
+ fn with_ballista_shuffle_reader_force_remote_read(
+ self,
+ force_remote_read: bool,
+ ) -> Self;
+
/// Sets parallelism of standalone cluster
///
/// This option to be used to configure standalone session context
@@ -319,6 +335,28 @@ impl SessionConfigExt for SessionConfig {
.set_usize(BALLISTA_SHUFFLE_READER_MAX_REQUESTS, max_requests)
}
}
+
+ fn ballista_shuffle_reader_force_remote_read(&self) -> bool {
+ self.options()
+ .extensions
+ .get::<BallistaConfig>()
+ .map(|c| c.shuffle_reader_force_remote_read())
+ .unwrap_or_else(|| {
+ BallistaConfig::default().shuffle_reader_force_remote_read()
+ })
+ }
+
+ fn with_ballista_shuffle_reader_force_remote_read(
+ self,
+ force_remote_read: bool,
+ ) -> Self {
+ if self.options().extensions.get::<BallistaConfig>().is_some() {
+ self.set_bool(BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ,
force_remote_read)
+ } else {
+ self.with_option_extension(BallistaConfig::default())
+ .set_bool(BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ,
force_remote_read)
+ }
+ }
}
impl SessionConfigHelperExt for SessionConfig {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
