1503: feat: (remote) shuffle reader cleanup - #6
Conversation
WalkthroughThis pull request refactors the Ballista client configuration and partition fetching mechanism. The ✨ Finishing Touches
🧪 Generate unit tests (beta)
📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request refactors and cleans up the remote shuffle reader client logic within Ballista. The primary goal is to centralize configuration management and simplify the API for fetching partitions. By moving connection details into the Highlights
Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request provides a good cleanup of the shuffle reader and related configuration handling. The BallistaClient is now more stateful, storing its host and port, which simplifies the fetch_partition API. Configuration handling is also improved by centralizing settings within BallistaConfig and removing separate wrapper structs, leading to more consistent code. The shuffle reader logic itself has been refactored for clarity, for example by removing the PartitionReader trait abstraction.
I have one suggestion regarding a performance optimization in the shuffle reader to avoid expensive cloning of BallistaConfig in a loop.
|
|
||
| spawned_tasks.push(SpawnedTask::spawn({ | ||
| // TODO: make BallistaConfig cheaper to clone | ||
| let config = config.clone(); |
There was a problem hiding this comment.
As noted in the TODO on the preceding line, cloning BallistaConfig here for every remote partition can be expensive due to the inner HashMap. This can be optimized by wrapping BallistaConfig in an Arc and cloning the Arc instead, which is a cheap reference count increment.
You could change send_fetch_partitions to accept an Arc<BallistaConfig>:
fn send_fetch_partitions(
partition_locations: Vec<PartitionLocation>,
config: Arc<BallistaConfig>,
customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
) -> AbortableReceiverStream { /* ... */ }Then, update the call sites to wrap the config in an Arc.
In ShuffleReaderExec::execute:
let config = Arc::new(context.session_config().ballista_config());
let response_receiver =
send_fetch_partitions(partition_locations, config, customize_endpoint);In tests::test_send_fetch_partitions:
let response_receiver = send_fetch_partitions(partition_locations, Arc::new(config), None);This would make the clone inside the spawned task a cheap Arc clone, addressing the performance concern.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
| DataFusionError::Internal("Received empty partition id".to_owned()) | ||
| })?; | ||
| let host = metadata.host.as_str(); | ||
| let port = metadata.port as u16; |
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
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.
|
PR Review: feat: (remote) shuffle reader cleanup Overall this is a solid refactoring PR. A few items to address before merging. ISSUES
BallistaClient::try_new is called with client_host/client_port, which may be a proxy address when a flight proxy is configured. Previously the host/port passed to fetch_partition came from metadata.host/metadata.port -- the actual executor address. Now self.host/self.port stored in the client are the connection address, and these get embedded in Action::FetchPartition. If the executor uses those fields to identify its own partitions, passing the proxy address instead of the executor address could be incorrect. Worth verifying this semantic difference is intentional or harmless. The shuffle reader path is unaffected since it connects directly to metadata.host/metadata.port.
The dead comment inside set_with_prefix should be removed. It also makes the method name misleading -- the method does not actually add any prefix. Renaming to set or set_config_value would be clearer.
A blank line is missing between get_f64_setting and the doc comment for set_with_prefix. Rust convention also starts doc comments with a capital letter. POSITIVES
MINOR NIT new_ballista_client is a thin wrapper around BallistaClient::try_new that only extracts two config values. Consider inlining it into fetch_partition_remote to reduce indirection, or keeping it as a seam for future connection pooling (the existing TODO comment suggests this is future work). |
🤖 Augment PR SummarySummary: Refactors remote shuffle partition reads and related client configuration to reduce duplicated plumbing and centralize settings. Changes:
Technical Notes: Remote shuffle reads now consistently derive concurrency/message-size/TLS/transport preferences from 🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| path: path.to_owned(), | ||
| host: host.to_owned(), | ||
| port, | ||
| host: self.host.to_owned(), |
There was a problem hiding this comment.
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
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
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.
| &location.path, | ||
| host, | ||
| port, | ||
| flight_transport, |
There was a problem hiding this comment.
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
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
There was a problem hiding this comment.
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
ballista/core/src/execution_plans/shuffle_reader.rs (1)
440-456: Config cloning per remote partition may be expensive at scale.The
config.clone()at line 442 occurs for each remote partition. For queries with many partitions, this could add overhead. The existing TODO acknowledges this.Consider extracting the needed config values (max_message_size, use_tls, prefer_flight) into a small struct before the loop to avoid repeated cloning:
struct RemoteFetchConfig { max_message_size: usize, use_tls: bool, prefer_flight: bool, }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@ballista/core/src/execution_plans/shuffle_reader.rs` around lines 440 - 456, The loop currently calls config.clone() for every spawned task which is expensive; extract only the needed fields (e.g., max_message_size, use_tls, prefer_flight) into a small Copy/Clone struct (e.g., RemoteFetchConfig) outside the loop and capture that in the SpawnedTask::spawn closure instead of cloning the full BallistaConfig; update the async block passed to SpawnedTask::spawn to use the new RemoteFetchConfig when calling fetch_partition_remote (or pass it through to a modified fetch_partition_remote signature) and remove the per-iteration config.clone() call so only the small struct is moved into each task.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@ballista/core/src/config.rs`:
- Around line 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.
---
Nitpick comments:
In `@ballista/core/src/execution_plans/shuffle_reader.rs`:
- Around line 440-456: The loop currently calls config.clone() for every spawned
task which is expensive; extract only the needed fields (e.g., max_message_size,
use_tls, prefer_flight) into a small Copy/Clone struct (e.g., RemoteFetchConfig)
outside the loop and capture that in the SpawnedTask::spawn closure instead of
cloning the full BallistaConfig; update the async block passed to
SpawnedTask::spawn to use the new RemoteFetchConfig when calling
fetch_partition_remote (or pass it through to a modified fetch_partition_remote
signature) and remove the per-iteration config.clone() call so only the small
struct is moved into each task.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: af13e352-29ea-407b-b25f-f9f7b2380187
📒 Files selected for processing (7)
ballista/core/src/client.rsballista/core/src/config.rsballista/core/src/execution_plans/distributed_query.rsballista/core/src/execution_plans/shuffle_reader.rsballista/core/src/extension.rsballista/core/src/utils.rsexamples/examples/standalone-substrait.rs
💤 Files with no reviewable changes (1)
- examples/examples/standalone-substrait.rs
| /// 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) | ||
| } | ||
| } |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
value:good-to-have; category:bug; feedback: The CodeRabbit AI is correct! Introducing RemoteFetchConfig will minimize the size of the copied/cloned fields. Another option is to use std::sync::Arc that uses reference counting. |
value:useful; category:bug; feedback: The Claude 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. |
2 participants
1503: To review by AI