andygrove commented on code in PR #2558: URL: https://github.com/apache/datafusion-ballista/pull/2558#discussion_r4213274978
########## ballista/scheduler/src/physical_optimizer/file_group_balance.rs: ########## @@ -0,0 +1,449 @@ +// 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. +//! +//! [`BalanceFileGroups`] 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 BalanceFileGroups; Review Comment: Good call, renamed it to `BalanceAndPruneFileGroups` in 168b80b19. -- 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]
