You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/10/08 12:34:07 UTC
(datafusion-ballista) branch main updated: fix(scheduler): prune scan files by statistics and balance file groups (#2558)
This is an automated email from the ASF dual-hosted git repository.
andygrove 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 cf2642fbe fix(scheduler): prune scan files by statistics and balance file groups (#2558)
cf2642fbe is described below
commit cf2642fbe938511e7c530515568a4b4c0eff1797
Author: Andy Grove <ag...@apache.org>
AuthorDate: Thu Oct 8 06:34:02 2026 -0600
fix(scheduler): prune scan files by statistics and balance file groups (#2558)
---
ballista/core/src/config.rs | 23 ++
.../src/physical_optimizer/file_group_balance.rs | 450 +++++++++++++++++++++
ballista/scheduler/src/physical_optimizer/mod.rs | 3 +
ballista/scheduler/src/state/aqe/planner.rs | 5 +
ballista/scheduler/src/state/task_manager.rs | 4 +
docs/source/user-guide/configs.md | 1 +
6 files changed, 486 insertions(+)
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index 31ea8ec0b..e0f9363ea 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -153,6 +153,13 @@ pub const BALLISTA_STAGE_BUILD_SIDE_MAX_ESTIMATE_MULTIPLE: &str =
pub const BALLISTA_NOT_IN_SUBQUERY_REWRITE: &str =
"ballista.optimizer.not_in_subquery_rewrite";
+/// Configuration key controlling whether scan files are pruned by statistics
+/// during planning and the surviving files are spread evenly over the scan's
+/// file groups. Enabled by default; set to `false` to leave file groups as
+/// DataFusion built them.
+pub const BALLISTA_BALANCE_SCAN_FILE_GROUPS: &str =
+ "ballista.optimizer.balance_scan_file_groups";
+
/// Configuration key to enable AQE coalesce-shuffle-partitions rule.
/// Disabled by default — opt in when the workload benefits from larger
/// downstream tasks more than from preserved parallelism.
@@ -362,6 +369,16 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
to keep the null-aware join and its single-task lowering.".to_string(),
DataType::Boolean,
Some("true".to_string())),
+ ConfigEntry::new(BALLISTA_BALANCE_SCAN_FILE_GROUPS.to_string(),
+ "Prunes scan files using their Parquet statistics while planning and \
+ deals the remaining files round-robin over the scan's file groups. \
+ DataFusion builds file groups before any predicate is applied, so when \
+ files are clustered by the filtered column (for example a date range \
+ over date-partitioned files) the files that survive pruning land in a \
+ few neighbouring groups and only those tasks have work to do. Set to \
+ false to keep the file groups as DataFusion built them.".to_string(),
+ DataType::Boolean,
+ Some("true".to_string())),
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,
@@ -806,6 +823,12 @@ impl BallistaConfig {
self.get_bool_setting(BALLISTA_NOT_IN_SUBQUERY_REWRITE)
}
+ /// Whether scan files are pruned by statistics during planning and the
+ /// surviving files are spread evenly over the scan's file groups.
+ pub fn balance_scan_file_groups_enabled(&self) -> bool {
+ self.get_bool_setting(BALLISTA_BALANCE_SCAN_FILE_GROUPS)
+ }
+
/// Returns whether the AQE coalesce-shuffle-partitions rule is enabled.
pub fn coalesce_enabled(&self) -> bool {
self.get_bool_setting(BALLISTA_COALESCE_ENABLED)
diff --git a/ballista/scheduler/src/physical_optimizer/file_group_balance.rs b/ballista/scheduler/src/physical_optimizer/file_group_balance.rs
new file mode 100644
index 000000000..ec6f146ee
--- /dev/null
+++ b/ballista/scheduler/src/physical_optimizer/file_group_balance.rs
@@ -0,0 +1,450 @@
+// 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.
+
+//! Prunes scan files by statistics at planning time and spreads the survivors
+//! evenly over the scan's file groups.
+//!
+//! DataFusion builds a file scan's file groups before any predicate is applied:
+//! the files are sorted by path and cut into contiguous runs. Statistics-based
+//! pruning only happens later, inside each task, as it opens its files. When
+//! the files are clustered by the filtered column (for example Hive-style
+//! `l_shipdate=...` directories filtered by a date range), every file that
+//! survives sits in a few neighbouring groups. Only the tasks that own those
+//! groups have rows to read, the rest open and skip their files, and the stage
+//! takes as long as the busiest few tasks.
+//!
+//! [`BalanceAndPruneFileGroups`] evaluates the scan's predicate against each
+//! file's statistics while planning, drops the files that cannot match, and
+//! deals the remaining files round-robin over the groups so the surviving work
+//! is spread across tasks. It leaves a scan untouched when no file can be
+//! pruned.
+
+use ballista_core::config::BallistaConfig;
+use datafusion::common::Result;
+use datafusion::common::config::ConfigOptions;
+use datafusion::common::pruning::PrunableStatistics;
+use datafusion::common::tree_node::{Transformed, TreeNode};
+use datafusion::datasource::listing::PartitionedFile;
+use datafusion::datasource::physical_plan::{
+ FileGroup, FileScanConfig, FileScanConfigBuilder,
+};
+use datafusion::datasource::source::DataSourceExec;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_optimizer::pruning::PruningPredicateBuilder;
+use datafusion::physical_plan::ExecutionPlan;
+use log::debug;
+use std::sync::Arc;
+
+/// Physical optimizer rule that prunes file scans by file statistics and
+/// rebalances the surviving files across the scan's file groups.
+///
+/// See the [module documentation](self) for the problem it solves.
+///
+/// A scan is left untouched when any of the following holds, because the
+/// layout of its groups carries meaning that must not be disturbed:
+///
+/// * it declares an output partitioning (files are grouped by partition value),
+/// * it declares an output ordering, or must preserve file order,
+/// * it has a single file group,
+/// * it has no predicate, or the predicate cannot prune anything,
+/// * no file can be pruned from the statistics collected during planning.
+///
+/// Files without statistics are always kept. The number of file groups never
+/// grows, and shrinks only when fewer files survive than there were groups.
+#[derive(Debug, Default)]
+pub struct BalanceAndPruneFileGroups;
+
+impl BalanceAndPruneFileGroups {
+ /// Creates a new `BalanceAndPruneFileGroups` rule.
+ pub fn new() -> Self {
+ Self
+ }
+}
+
+impl PhysicalOptimizerRule for BalanceAndPruneFileGroups {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // On unless a Ballista session turns it off.
+ let enabled = config
+ .extensions
+ .get::<BallistaConfig>()
+ .is_none_or(|c| c.balance_scan_file_groups_enabled());
+ if !enabled {
+ return Ok(plan);
+ }
+
+ plan.transform_up(|plan| {
+ let Some(exec) = plan.downcast_ref::<DataSourceExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+ let Some(config) = exec.data_source().downcast_ref::<FileScanConfig>() else {
+ return Ok(Transformed::no(plan));
+ };
+ let Some(file_groups) = balance_file_groups(config) else {
+ return Ok(Transformed::no(plan));
+ };
+ let config = FileScanConfigBuilder::from(config.clone())
+ .with_file_groups(file_groups)
+ .build();
+ Ok(Transformed::yes(DataSourceExec::from_data_source(config)))
+ })
+ .map(|t| t.data)
+ }
+
+ fn name(&self) -> &str {
+ "BalanceAndPruneFileGroups"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Returns the rebalanced file groups for `config`, or `None` when the scan
+/// should keep its current groups.
+fn balance_file_groups(config: &FileScanConfig) -> Option<Vec<FileGroup>> {
+ let group_count = config.file_groups.len();
+ if group_count <= 1
+ || config.output_partitioning.is_some()
+ || config.preserve_order
+ || !config.output_ordering.is_empty()
+ {
+ return None;
+ }
+
+ let predicate = config.file_source().filter()?;
+ // `build` yields `None` when the predicate cannot prune anything.
+ let pruning_predicate = PruningPredicateBuilder::new()
+ .with_file_schema(Arc::clone(config.file_schema()))
+ .build(predicate)?;
+
+ // Each file is judged on its own, because one file with missing or inexact
+ // statistics would otherwise stop the whole batch from pruning anything.
+ // A file that cannot be judged is kept.
+ let keep = |file: &PartitionedFile| -> bool {
+ let Some(stats) = &file.statistics else {
+ return true;
+ };
+ let prunable = PrunableStatistics::new(
+ vec![Arc::clone(stats)],
+ Arc::clone(config.file_schema()),
+ );
+ pruning_predicate
+ .prune(&prunable)
+ .ok()
+ .is_none_or(|keep| keep[0])
+ };
+
+ let total: usize = config.file_groups.iter().map(|group| group.len()).sum();
+ let survivors: Vec<&PartitionedFile> = config
+ .file_groups
+ .iter()
+ .flat_map(|group| group.iter())
+ .filter(|file| keep(file))
+ .collect();
+ if survivors.len() == total {
+ return None;
+ }
+ debug!(
+ "Pruned {} scan files by statistics, {} remain over up to {group_count} file groups",
+ total - survivors.len(),
+ survivors.len()
+ );
+
+ // Never more groups than files, so no task is handed an empty group, and
+ // never fewer than one, so the scan keeps a partition to run.
+ let new_group_count = survivors.len().clamp(1, group_count);
+ let mut groups = vec![FileGroup::default(); new_group_count];
+ for (idx, file) in survivors.into_iter().enumerate() {
+ groups[idx % new_group_count].push(file.clone());
+ }
+ Some(groups)
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use ballista_core::config::BALLISTA_BALANCE_SCAN_FILE_GROUPS;
+ use ballista_core::extension::SessionConfigExt;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
+ use datafusion::common::stats::Precision;
+ use datafusion::common::{ColumnStatistics, ScalarValue, Statistics};
+ use datafusion::datasource::physical_plan::ParquetSource;
+ use datafusion::execution::object_store::ObjectStoreUrl;
+ use datafusion::logical_expr::Operator;
+ use datafusion::physical_expr::PhysicalExpr;
+ use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal};
+ use datafusion::physical_plan::Partitioning;
+ use datafusion::prelude::SessionConfig;
+
+ fn schema() -> SchemaRef {
+ Arc::new(Schema::new(vec![Field::new("d", DataType::Int64, false)]))
+ }
+
+ /// A file whose column `d` spans exactly `min..=max`.
+ fn file(name: &str, min: i64, max: i64) -> PartitionedFile {
+ let stats = Statistics {
+ num_rows: Precision::Exact(10),
+ total_byte_size: Precision::Exact(100),
+ column_statistics: vec![ColumnStatistics {
+ null_count: Precision::Exact(0),
+ min_value: Precision::Exact(ScalarValue::Int64(Some(min))),
+ max_value: Precision::Exact(ScalarValue::Int64(Some(max))),
+ ..ColumnStatistics::new_unknown()
+ }],
+ };
+ let mut file = PartitionedFile::new(name.to_string(), 100);
+ file.statistics = Some(Arc::new(stats));
+ file
+ }
+
+ /// `d >= lo AND d < hi`
+ fn range_predicate(lo: i64, hi: i64) -> Arc<dyn PhysicalExpr> {
+ let col = || Arc::new(Column::new("d", 0)) as Arc<dyn PhysicalExpr>;
+ let lit = |v: i64| {
+ Arc::new(Literal::new(ScalarValue::Int64(Some(v)))) as Arc<dyn PhysicalExpr>
+ };
+ Arc::new(BinaryExpr::new(
+ Arc::new(BinaryExpr::new(col(), Operator::GtEq, lit(lo))),
+ Operator::And,
+ Arc::new(BinaryExpr::new(col(), Operator::Lt, lit(hi))),
+ ))
+ }
+
+ /// A scan over `groups`, filtered by `predicate` when given.
+ fn scan(
+ groups: Vec<Vec<PartitionedFile>>,
+ predicate: Option<Arc<dyn PhysicalExpr>>,
+ ) -> FileScanConfigBuilder {
+ let mut source = ParquetSource::new(schema());
+ if let Some(predicate) = predicate {
+ source = source.with_predicate(predicate);
+ }
+ let mut builder = FileScanConfigBuilder::new(
+ ObjectStoreUrl::local_filesystem(),
+ Arc::new(source),
+ );
+ for group in groups {
+ builder = builder.with_file_group(FileGroup::new(group));
+ }
+ builder
+ }
+
+ fn optimize(builder: FileScanConfigBuilder) -> Arc<dyn ExecutionPlan> {
+ let plan = DataSourceExec::from_data_source(builder.build());
+ BalanceAndPruneFileGroups::new()
+ .optimize(plan, &ConfigOptions::default())
+ .unwrap()
+ }
+
+ fn scan_config(plan: &Arc<dyn ExecutionPlan>) -> &FileScanConfig {
+ let exec = plan.downcast_ref::<DataSourceExec>().unwrap();
+ exec.data_source().downcast_ref::<FileScanConfig>().unwrap()
+ }
+
+ fn group_names(plan: &Arc<dyn ExecutionPlan>) -> Vec<Vec<String>> {
+ scan_config(plan)
+ .file_groups
+ .iter()
+ .map(|g| g.iter().map(|f| f.path().to_string()).collect())
+ .collect()
+ }
+
+ /// Asserts the scan still has the four contiguous groups of
+ /// `clustered_groups`.
+ fn assert_unchanged(plan: &Arc<dyn ExecutionPlan>) {
+ assert_eq!(
+ group_names(plan),
+ vec![
+ vec!["f0", "f1"],
+ vec!["f2", "f3"],
+ vec!["f4", "f5"],
+ vec!["f6", "f7"]
+ ]
+ );
+ }
+
+ /// Four contiguous groups of two files each; only the files of group 1
+ /// and the first file of group 2 are in range.
+ fn clustered_groups() -> Vec<Vec<PartitionedFile>> {
+ (0..4)
+ .map(|g| {
+ (0..2)
+ .map(|i| {
+ let n = g * 2 + i;
+ file(&format!("f{n}"), n * 10, n * 10 + 9)
+ })
+ .collect()
+ })
+ .collect()
+ }
+
+ #[test]
+ fn spreads_surviving_files_across_groups() {
+ // Files f2..f5 (min 20..50) survive d in [20, 60).
+ let plan = optimize(scan(clustered_groups(), Some(range_predicate(20, 60))));
+ assert_eq!(
+ group_names(&plan),
+ vec![vec!["f2"], vec!["f3"], vec!["f4"], vec!["f5"]]
+ );
+ assert_eq!(plan.properties().output_partitioning().partition_count(), 4);
+ }
+
+ #[test]
+ fn shrinks_to_the_number_of_surviving_files() {
+ let plan = optimize(scan(clustered_groups(), Some(range_predicate(20, 40))));
+ assert_eq!(group_names(&plan), vec![vec!["f2"], vec!["f3"]]);
+ assert_eq!(plan.properties().output_partitioning().partition_count(), 2);
+ }
+
+ #[test]
+ fn keeps_one_group_when_every_file_is_pruned() {
+ let plan = optimize(scan(clustered_groups(), Some(range_predicate(500, 600))));
+ assert_eq!(group_names(&plan), vec![Vec::<String>::new()]);
+ }
+
+ #[test]
+ fn deals_surplus_files_round_robin() {
+ // 6 survivors (f2..f7) over 4 groups.
+ let plan = optimize(scan(clustered_groups(), Some(range_predicate(20, 80))));
+ assert_eq!(
+ group_names(&plan),
+ vec![vec!["f2", "f6"], vec!["f3", "f7"], vec!["f4"], vec!["f5"]]
+ );
+ }
+
+ #[test]
+ fn leaves_scan_alone_when_nothing_is_pruned() {
+ let plan = optimize(scan(clustered_groups(), Some(range_predicate(0, 1000))));
+ assert_unchanged(&plan);
+ }
+
+ #[test]
+ fn leaves_scan_alone_without_a_predicate() {
+ assert_unchanged(&optimize(scan(clustered_groups(), None)));
+ }
+
+ #[test]
+ fn prunes_around_files_without_usable_statistics() {
+ // `nostats` has no statistics at all, `partial` only an inexact
+ // minimum. Neither can be judged, so both are kept.
+ let mut groups = clustered_groups();
+ groups[0][0] = PartitionedFile::new("nostats".to_string(), 100);
+ let mut partial = file("partial", 0, 9);
+ let mut stats = Statistics::clone(partial.statistics.as_ref().unwrap());
+ stats.column_statistics[0].min_value =
+ Precision::Inexact(ScalarValue::Int64(Some(0)));
+ stats.column_statistics[0].max_value = Precision::Absent;
+ partial.statistics = Some(Arc::new(stats));
+ groups[0][1] = partial;
+
+ let plan = optimize(scan(groups, Some(range_predicate(20, 40))));
+ assert_eq!(
+ group_names(&plan),
+ vec![vec!["nostats"], vec!["partial"], vec!["f2"], vec!["f3"]]
+ );
+ }
+
+ #[test]
+ fn leaves_scan_alone_with_declared_output_partitioning() {
+ let key: Arc<dyn PhysicalExpr> = Arc::new(Column::new("d", 0));
+ let builder = scan(clustered_groups(), Some(range_predicate(20, 60)))
+ .with_output_partitioning(Some(Partitioning::Hash(vec![key], 4)));
+ assert_unchanged(&optimize(builder));
+ }
+
+ #[test]
+ fn leaves_scan_alone_when_order_must_be_preserved() {
+ let builder = scan(clustered_groups(), Some(range_predicate(20, 60)))
+ .with_preserve_order(true);
+ assert_unchanged(&optimize(builder));
+ }
+
+ #[test]
+ fn leaves_scan_alone_when_disabled() -> Result<()> {
+ let config = SessionConfig::new_with_ballista()
+ .set_bool(BALLISTA_BALANCE_SCAN_FILE_GROUPS, false);
+ let plan = DataSourceExec::from_data_source(
+ scan(clustered_groups(), Some(range_predicate(20, 60))).build(),
+ );
+ let plan = BalanceAndPruneFileGroups::new().optimize(plan, config.options())?;
+ assert_unchanged(&plan);
+ Ok(())
+ }
+
+ /// Plans a real `ListingTable` over 16 Parquet files clustered by `d`, so
+ /// the rule is exercised against the statistics and pushed-down predicate
+ /// DataFusion actually produces.
+ #[tokio::test]
+ async fn balances_a_planned_listing_table() -> Result<()> {
+ use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch};
+ use datafusion::parquet::arrow::ArrowWriter;
+ use datafusion::prelude::{ParquetReadOptions, SessionContext};
+
+ let dir = tempfile::tempdir()?;
+ for n in 0..16i64 {
+ let values: ArrayRef =
+ Arc::new(Int64Array::from_iter_values((n * 10)..(n * 10 + 10)));
+ let batch = RecordBatch::try_new(schema(), vec![values])?;
+ let path = dir.path().join(format!("f{n:02}.parquet"));
+ let mut writer =
+ ArrowWriter::try_new(std::fs::File::create(path)?, schema(), None)?;
+ writer.write(&batch)?;
+ writer.close()?;
+ }
+
+ let ctx = SessionContext::new_with_config(
+ SessionConfig::new().with_target_partitions(4),
+ );
+ ctx.register_parquet(
+ "t",
+ dir.path().to_str().unwrap(),
+ ParquetReadOptions::default(),
+ )
+ .await?;
+ // Files f04..f07 hold d in [40, 80).
+ let plan = ctx
+ .sql("SELECT d FROM t WHERE d >= 40 AND d < 80")
+ .await?
+ .create_physical_plan()
+ .await?;
+
+ let scan_groups = |plan: &Arc<dyn ExecutionPlan>| -> Vec<usize> {
+ let mut leaf = plan;
+ while let Some(child) = leaf.children().first() {
+ leaf = child;
+ }
+ scan_config(leaf)
+ .file_groups
+ .iter()
+ .map(|g| g.len())
+ .collect()
+ };
+
+ // DataFusion cuts 16 files into 4 contiguous runs of 4.
+ assert_eq!(scan_groups(&plan), vec![4, 4, 4, 4]);
+
+ let balanced =
+ BalanceAndPruneFileGroups::new().optimize(plan, &ConfigOptions::default())?;
+ // Only 4 files survive, one per group.
+ assert_eq!(scan_groups(&balanced), vec![1, 1, 1, 1]);
+ Ok(())
+ }
+}
diff --git a/ballista/scheduler/src/physical_optimizer/mod.rs b/ballista/scheduler/src/physical_optimizer/mod.rs
index 4252d9549..56b700d71 100644
--- a/ballista/scheduler/src/physical_optimizer/mod.rs
+++ b/ballista/scheduler/src/physical_optimizer/mod.rs
@@ -25,3 +25,6 @@ pub mod filter_pushdown;
// specific cases. it has been used in static
// execution graph only.
pub mod join_selection;
+// prunes scan files by statistics at planning time
+// and spreads the survivors over the scan's file groups
+pub mod file_group_balance;
diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs
index 2a01df0f9..dd4ac69cf 100644
--- a/ballista/scheduler/src/state/aqe/planner.rs
+++ b/ballista/scheduler/src/state/aqe/planner.rs
@@ -14,6 +14,7 @@
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
+use crate::physical_optimizer::file_group_balance::BalanceAndPruneFileGroups;
use crate::physical_optimizer::filter_pushdown::FilterPushdown;
use crate::state::aqe::adapter::BallistaAdapter;
use crate::state::aqe::execution_plan::{
@@ -625,6 +626,10 @@ impl AdaptivePlanner {
) -> Vec<PhysicalOptimizerRuleRef> {
vec![
Arc::new(FilterPushdown::new()),
+ // Needs the predicate `FilterPushdown` leaves on each scan. Runs
+ // before DataFusion's own chain so its file-scan repartitioning
+ // sees the pruned files.
+ Arc::new(BalanceAndPruneFileGroups::new()),
Arc::new(DelayJoinSelectionRule::new(plan_id_generator)),
Arc::new(ChaosCreatingRule::default()),
]
diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs
index 1cae06582..621316d20 100644
--- a/ballista/scheduler/src/state/task_manager.rs
+++ b/ballista/scheduler/src/state/task_manager.rs
@@ -17,6 +17,7 @@
use crate::cluster::{ExecutorSlot, JobState, JobStateEventStream};
use crate::config::SchedulerConfig;
+use crate::physical_optimizer::file_group_balance::BalanceAndPruneFileGroups;
use crate::planner::DefaultDistributedPlanner;
use crate::scheduler_server::event::{QueryStageSchedulerEvent, SubmitPlan};
use crate::state::aqe::AdaptiveExecutionGraph;
@@ -40,6 +41,7 @@ use dashmap::DashMap;
use datafusion::execution::config::SessionConfig;
use datafusion::execution::context::SessionContext;
use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::physical_plan::ExecutionPlan;
use datafusion_proto::logical_plan::AsLogicalPlan;
use datafusion_proto::physical_plan::AsExecutionPlan;
@@ -291,6 +293,8 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
let physical_plan =
ctx.state().create_physical_plan(logical_plan).await?;
+ let physical_plan = BalanceAndPruneFileGroups::new()
+ .optimize(physical_plan, session_config.options())?;
let physical_plan =
handle_explain_plan(job_id, &ctx, logical_plan, physical_plan)
.await?;
diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md
index 9f3a36732..9df1614b1 100644
--- a/docs/source/user-guide/configs.md
+++ b/docs/source/user-guide/configs.md
@@ -119,6 +119,7 @@ standard DataFusion settings.
| ballista.client.pull | Boolean | false | 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. [...]
| ballista.client.use_tls | Boolean | false | Should connection between client, scheduler, and executors use TLS. [...]
| ballista.job.name | Utf8 | (none) | Sets the job name that will appear in the web user interface for any submitted jobs [...]
+| ballista.optimizer.balance_scan_file_groups | Boolean | true | Prunes scan files using their Parquet statistics while planning and deals the remaining files round-robin over the scan's file groups. DataFusion builds file groups before any predicate is applied, so when files are clustered by the filtered column (for example a date range over date-partitioned files) the files that survive pruning land in a few neighbouring groups and only those task [...]
| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 134217728 | Byte-size threshold below which a hash join's smaller side is promoted to CollectLeft and lowered via the broadcast pattern. Governs broadcast selection under both the static distributed planner and adaptive query planning (AQE). It also caps null-aware anti joins with a known build size because they require single-task CollectLeft execution. Set to 0 to disable promotion and reject nu [...]
| ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 | Row-count threshold below which a hash join's smaller side is promoted to CollectLeft and lowered via the broadcast pattern, used as a fallback when byte-size statistics are unavailable. Applies to adaptive query planning (AQE). Set to 0 to disable promotion via the row-count path. [...]
| ballista.optimizer.hash_join_max_build_partition_bytes | UInt64 | 67108864 | Maximum per-partition hash-join build-side bytes for a Partitioned hash join under AQE. A build partition larger than this falls back to SortMergeJoin (spillable). Defaults to 64 MiB; 0 disables the check, which makes AQE use a hash join regardless of build size. [...]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org