Skip to content

1503: feat: (remote) shuffle reader cleanup #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
martin-augment wants to merge 1 commit into
base: main
Choose a base branch
from
Open

1503: feat: (remote) shuffle reader cleanup #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 18 additions & 20 deletions ballista/core/src/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,11 @@

//! Client API for sending requests to executors.

use std::collections::HashMap;
use std::sync::Arc;

use std::{
convert::{TryFrom, TryInto},
task::{Context, Poll},
};

use crate::error::{BallistaError, Result as BResult};
use crate::extension::BallistaConfigGrpcEndpoint;
use crate::serde::protobuf;
use crate::serde::scheduler::{Action, PartitionId};

use crate::utils::create_grpc_client_endpoint;
use arrow_flight;
use arrow_flight::Ticket;
use arrow_flight::utils::flight_data_to_arrow_batch;
Expand All @@ -43,21 +37,23 @@ use datafusion::arrow::{
};
use datafusion::error::DataFusionError;
use datafusion::error::Result;

use crate::extension::BallistaConfigGrpcEndpoint;
use crate::serde::protobuf;

use crate::utils::create_grpc_client_endpoint;

use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream};
use futures::{Stream, StreamExt};
use log::{debug, warn};
use prost::Message;
use std::collections::HashMap;
use std::sync::Arc;
use std::{
convert::{TryFrom, TryInto},
task::{Context, Poll},
};
use tonic::{Code, Streaming};

/// Client for interacting with Ballista executors.
#[derive(Clone)]
pub struct BallistaClient {
host: String,
port: u16,
flight_client: FlightServiceClient<tonic::transport::channel::Channel>,
}

Expand Down Expand Up @@ -109,7 +105,11 @@ impl BallistaClient {

debug!("BallistaClient connected OK: {flight_client:?}");

Ok(Self { flight_client })
Ok(Self {
flight_client,
host: host.to_string(),
port,
})
}

/// Retrieves a partition from an executor.
Expand All @@ -122,17 +122,15 @@ impl BallistaClient {
executor_id: &str,
partition_id: &PartitionId,
path: &str,
host: &str,
port: u16,
flight_transport: bool,
) -> BResult<SendableRecordBatchStream> {
let action = Action::FetchPartition {
job_id: partition_id.job_id.clone(),
stage_id: partition_id.stage_id,
partition_id: partition_id.partition_id,
path: path.to_owned(),
host: host.to_owned(),
port,
host: self.host.to_owned(),

@augmentcode augmentcode Bot Mar 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FetchPartition.host/port are used by the scheduler flight proxy (flight_proxy_service.rs) to connect onward to the executor; using self.host/self.port here looks problematic when the client is connected to a proxy/scheduler endpoint. In those cases the ticket may embed the proxy’s host/port and cause the proxy to forward back to itself rather than the executor.

Severity: high

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

@martin-augment martin-augment Mar 13, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:useful; category:bug; feedback: The Augment AI reviewer is correct! The changes in the Pull Request propose to use the flight proxy for some functionalities and still use direct connection to the executor for others. In this particular case the connection should be direct/sticky, because a specific executor is responsible for a given partition.

port: self.port,
};

let result = if flight_transport {
Expand Down
51 changes: 38 additions & 13 deletions ballista/core/src/config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,8 @@ use datafusion::{
pub const BALLISTA_JOB_NAME: &str = "ballista.job.name";
/// Configuration key for standalone processing parallelism.
pub const BALLISTA_STANDALONE_PARALLELISM: &str = "ballista.standalone.parallelism";

/// Configuration key for disabling default cache extension node.
pub const BALLISTA_CACHE_NOOP: &str = "ballista.cache.noop";

/// Configuration key for maximum concurrent shuffle read requests.
pub const BALLISTA_SHUFFLE_READER_MAX_REQUESTS: &str =
"ballista.shuffle.max_concurrent_read_requests";
Expand All @@ -44,7 +42,6 @@ pub const BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ: &str =
/// Configuration key to prefer Flight protocol for remote shuffle reads.
pub const BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT: &str =
"ballista.shuffle.remote_read_prefer_flight";

/// max message size for gRPC clients
pub const BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE: &str =
"ballista.grpc_client_max_message_size";
Expand Down Expand Up @@ -82,6 +79,8 @@ pub const BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE: &str =
"ballista.shuffle.sort_based.batch_size";
/// Should client employ pull or push job tracking strategy
pub const BALLISTA_CLIENT_PULL: &str = "ballista.client.pull";
/// Should client use tls connection
pub const BALLISTA_CLIENT_USE_TLS: &str = "ballista.client.use_tls";

/// Result type for configuration parsing operations.
pub type ParseResult<T> = result::Result<T, String>;
Expand Down Expand Up @@ -162,6 +161,10 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
ConfigEntry::new(BALLISTA_CLIENT_PULL.to_string(),
"Should client employ pull or push job tracking. In pull mode client will make a request to server in the loop, until job finishes. Pull mode is kept for legacy clients.".to_string(),
DataType::Boolean,
Some(false.to_string())),
ConfigEntry::new(BALLISTA_CLIENT_USE_TLS.to_string(),
"Should connection between client, scheduler, and executors use TLS.".to_string(),
DataType::Boolean,
Some(false.to_string()))
];
entries
Expand Down Expand Up @@ -274,11 +277,6 @@ impl BallistaConfig {
&self.settings
}

/// Returns the maximum message size for gRPC clients in bytes.
pub fn default_grpc_client_max_message_size(&self) -> usize {
self.get_usize_setting(BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE)
}

/// Returns the standalone processing parallelism level.
pub fn default_standalone_parallelism(&self) -> usize {
self.get_usize_setting(BALLISTA_STANDALONE_PARALLELISM)
Expand All @@ -290,25 +288,30 @@ impl BallistaConfig {
}

/// Returns the gRPC client connection timeout in seconds.
pub fn default_grpc_client_connect_timeout_seconds(&self) -> usize {
pub fn grpc_client_connect_timeout_seconds(&self) -> usize {
self.get_usize_setting(BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS)
}

/// Returns the gRPC client request timeout in seconds.
pub fn default_grpc_client_timeout_seconds(&self) -> usize {
pub fn grpc_client_timeout_seconds(&self) -> usize {
self.get_usize_setting(BALLISTA_GRPC_CLIENT_TIMEOUT_SECONDS)
}

/// Returns the TCP keep-alive interval for gRPC clients in seconds.
pub fn default_grpc_client_tcp_keepalive_seconds(&self) -> usize {
pub fn grpc_client_tcp_keepalive_seconds(&self) -> usize {
self.get_usize_setting(BALLISTA_GRPC_CLIENT_TCP_KEEPALIVE_SECONDS)
}

/// Returns the HTTP/2 keep-alive interval for gRPC clients in seconds.
pub fn default_grpc_client_http2_keepalive_interval_seconds(&self) -> usize {
pub fn grpc_client_http2_keepalive_interval_seconds(&self) -> usize {
self.get_usize_setting(BALLISTA_GRPC_CLIENT_HTTP2_KEEPALIVE_INTERVAL_SECONDS)
}

/// Returns the maximum message size for gRPC clients in bytes.
pub fn grpc_client_max_message_size(&self) -> usize {
self.get_usize_setting(BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE)
}

/// Returns whether the default cache node extension is disabled.
pub fn cache_noop(&self) -> bool {
self.get_bool_setting(BALLISTA_CACHE_NOOP)
Expand Down Expand Up @@ -373,6 +376,11 @@ impl BallistaConfig {
self.get_bool_setting(BALLISTA_CLIENT_PULL)
}

/// should client use TLS to communicate with ballista cluster
pub fn client_use_tls(&self) -> bool {
self.get_bool_setting(BALLISTA_CLIENT_USE_TLS)
}

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
Expand Down Expand Up @@ -419,6 +427,23 @@ impl BallistaConfig {
v.parse::<f64>().unwrap()
}
}
/// sets the configuration value where key starts with ballista
/// prefix.
pub fn set_with_prefix(
&mut self,
key: &str,
value: &str,
) -> datafusion::error::Result<()> {
let entries = Self::valid_entries();
//let k = format!("{}.{key}", BallistaConfig::PREFIX);

if entries.contains_key(key) {
self.settings.insert(key.to_string(), value.to_string());
Ok(())
} else {
config_err!("configuration key `{}` does not exist", key)
}
}
Comment on lines +430 to +446

@coderabbitai coderabbitai Bot Mar 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

set_with_prefix does not validate value types.

Unlike with_settings() which validates that values can be parsed as their declared types, set_with_prefix only validates key existence. Setting an invalid value (e.g., "not_a_number" for a UInt64 config) will cause a panic when the value is later accessed via get_usize_setting() which calls .unwrap().

Consider adding type validation:

🛡️ Proposed fix to add value validation 
 pub fn set_with_prefix(
     &mut self,
     key: &str,
     value: &str,
 ) -> datafusion::error::Result<()> {
     let entries = Self::valid_entries();
 
     if entries.contains_key(key) {
+        let entry = entries.get(key).unwrap();
+        Self::parse_value(value, entry.data_type.clone())
+            .map_err(|e| datafusion::error::DataFusionError::Configuration(
+                format!("Invalid value '{value}' for key '{key}': {e}")
+            ))?;
         self.settings.insert(key.to_string(), value.to_string());
         Ok(())
     } else {
         config_err!("configuration key `{}` does not exist", key)
     }
 }
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@ballista/core/src/config.rs` around lines 430 - 446, The set_with_prefix
function currently only checks key existence (use Self::valid_entries and
entries.contains_key) but does not validate that value parses to the declared
type, which later causes panics in callers like get_usize_setting that unwrap
parsed values; update set_with_prefix to fetch the entry from valid_entries
(e.g., entries.get(key)), attempt to parse the provided value into the entry's
declared type using the same parsing/validation logic used by with_settings(),
and only insert into self.settings if parsing succeeds; on parse failure return
a config_err! with a clear message rather than inserting the invalid string.
Ensure you reference and reuse the same parsing helper/path used by
with_settings() to keep behavior consistent.

@martin-augment martin-augment Mar 13, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:useful; category:bug; feedback: The CodeRabbit AI reviewer is correct! The new method does not validate the provided values and this may lead to problems later when these values need to be used. Prevents storing invalid values.

}

impl datafusion::config::ExtensionOptions for BallistaConfig {
Expand Down Expand Up @@ -539,7 +564,7 @@ mod tests {
#[test]
fn default_config() -> Result<()> {
let config = BallistaConfig::default();
assert_eq!(16777216, config.default_grpc_client_max_message_size());
assert_eq!(16777216, config.grpc_client_max_message_size());
Ok(())
}
}
8 changes: 2 additions & 6 deletions ballista/core/src/execution_plans/distributed_query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
self.scheduler_url.clone(),
self.session_id.clone(),
query,
self.config.default_grpc_client_max_message_size(),
self.config.grpc_client_max_message_size(),
GrpcClientConfig::from(&self.config),
Arc::new(self.metrics.clone()),
partition,
Expand All @@ -280,7 +280,7 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
execute_query_push(
self.scheduler_url.clone(),
query,
self.config.default_grpc_client_max_message_size(),
self.config.grpc_client_max_message_size(),
GrpcClientConfig::from(&self.config),
Arc::new(self.metrics.clone()),
partition,
Expand Down Expand Up @@ -701,8 +701,6 @@ async fn fetch_partition(
let partition_id = location.partition_id.ok_or_else(|| {
DataFusionError::Internal("Received empty partition id".to_owned())
})?;
let host = metadata.host.as_str();
let port = metadata.port as u16;

@cursor cursor Bot Mar 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Action contains proxy address instead of executor address

Medium Severity

When a flight proxy is configured, get_client_host_port returns the proxy's address (scheduler or external proxy), not the executor's. The old code explicitly passed metadata.host and metadata.port to fetch_partition for the Action::FetchPartition payload, keeping the executor's real address. Now BallistaClient stores the connection target (client_host/client_port) as self.host/self.port, and fetch_partition uses those in the action. This means the Action::FetchPartition host/port fields contain the proxy address instead of the executor's actual address in proxy scenarios.

Additional Locations (1)
Fix in Cursor Fix in Web

@martin-augment martin-augment Mar 13, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:useful; category:bug; feedback: The Bugbot AI reviewer is correct! The changes in the Pull Request propose to use the flight proxy for some functionalities and still use direct connection to the executor for others. In this particular case the connection should be direct/sticky, because a specific executor is responsible for a given partition.


let (client_host, client_port) =
get_client_host_port(&metadata, &scheduler_url, &flight_proxy)?;
Expand All @@ -721,8 +719,6 @@ async fn fetch_partition(
&metadata.id,
&partition_id.into(),
&location.path,
host,
port,
flight_transport,

@augmentcode augmentcode Bot Mar 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In FlightProxy::Local(true) / FlightProxy::External modes get_client_host_port makes the Flight client connect to the proxy/scheduler, but the FetchPartition ticket still needs the executor’s metadata.host/port for proxy routing. With the BallistaClient::fetch_partition signature change, this call path seems likely to produce tickets that point back at the proxy/scheduler (potential forwarding loop/failure).

Severity: high

Other Locations
  • examples/examples/standalone-substrait.rs:419

Fix This in Augment

🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.

@martin-augment martin-augment Mar 13, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

value:useful; category:bug; feedback: The Augment AI reviewer is correct! The changes in the Pull Request propose to use the flight proxy for some functionalities and still use direct connection to the executor for others. In this particular case the connection should be direct/sticky, because a specific executor is responsible for a given partition.

)
.await
Expand Down
Loading
Loading