feat: Prune ListingTable file groups for range-key filters

datafusion/catalog-listing/Cargo.toml

@@ -42,6 +42,7 @@ datafusion-physical-expr = { workspace = true }
 datafusion-physical-expr-adapter = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-plan = { workspace = true }
+datafusion-pruning = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }

datafusion/catalog-listing/src/table.rs

@@ -45,15 +45,19 @@ use datafusion_execution::cache::cache_manager::{
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
+use datafusion_expr::utils::conjunction;
 use datafusion_expr::{
     Expr, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType,
 };
-use datafusion_physical_expr::{create_lex_ordering, create_physical_partitioning};
+use datafusion_physical_expr::{
+    Partitioning, create_lex_ordering, create_physical_expr, create_physical_partitioning,
+};
 use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::empty::EmptyExec;
 use datafusion_physical_plan::limit::GlobalLimitExec;
+use datafusion_pruning::{PruningPredicateBuilder, RangePruningStatistics};
 use futures::future::BoxFuture;
 use futures::{Stream, StreamExt, TryStreamExt, future, stream};
 use object_store::ObjectStore;
@@ -647,10 +651,17 @@ impl ListingTable {
 
         // If the filters can be resolved using only partition cols, there is no need to
         // pushdown it to TableScan, otherwise, `unhandled` pruning predicates will be generated
-        let (partition_filters, filters): (Vec<_>, Vec<_>) =
-            filters.iter().cloned().partition(|filter| {
-                can_be_evaluated_for_partition_pruning(&table_partition_col_names, filter)
-            });
+        let scan_filters = filters;
+        let filters = scan_filters
+            .iter()
+            .filter(|filter| {
+                !can_be_evaluated_for_partition_pruning(
+                    &table_partition_col_names,
+                    filter,
+                )
+            })
+            .cloned()
+            .collect::<Vec<_>>();
 
         let declared_output_partitioning = self.options.output_partitioning.as_ref();
 
@@ -672,7 +683,7 @@ impl ListingTable {
             grouped_by_partition: partitioned_by_file_group,
             partition_split_points,
         } = self
-            .list_files_for_scan(state, &partition_filters, statistic_file_limit)
+            .list_files_for_scan(state, &scan_filters, statistic_file_limit)
             .await?;
 
         // if no files need to be read, return an `EmptyExec`
@@ -918,22 +929,41 @@ impl ListingTable {
     /// be distributed to different threads / executors.
     ///
     /// If [`ListingOptions::output_partitioning`] is set, returns one file
-    /// group per declared partition, including empty trailing groups.
+    /// group per declared partition, including empty trailing groups. Filters
+    /// known not to match declared ranges leave empty groups at those indices.
     pub async fn list_files_for_scan<'a>(
         &'a self,
         ctx: &'a dyn Session,
         filters: &'a [Expr],
         limit: Option<usize>,
     ) -> datafusion_common::Result<ListFilesResult> {
+        let partition_column_names = self
+            .options
+            .table_partition_cols
+            .iter()
+            .map(|(name, _)| name.as_str())
+            .collect::<Vec<_>>();
+        let partition_filters = filters
+            .iter()
+            .filter(|filter| {
+                can_be_evaluated_for_partition_pruning(&partition_column_names, filter)
+            })
+            .cloned()
+            .collect::<Vec<_>>();
         if let Some(output_partitioning) = self.options.output_partitioning.as_ref() {
             self.list_files_for_declared_output_partitioning(
                 ctx,
                 output_partitioning,
+                &partition_filters,
                 filters,
             )
             .await
         } else {
-            self.list_files_for_regular_scan(ctx, filters, limit).await
+            let file_limit = (partition_filters.len() == filters.len())
+                .then_some(limit)
+                .flatten();
+            self.list_files_for_regular_scan(ctx, &partition_filters, file_limit)
+                .await
         }
     }
 
@@ -1056,6 +1086,7 @@ impl ListingTable {
         ctx: &'a dyn Session,
         output_partitioning: &LogicalPartitioning,
         filters: &'a [Expr],
+        range_filters: &[Expr],
     ) -> datafusion_common::Result<ListFilesResult> {
         let Some(file_group_count) = output_partitioning.partition_count() else {
             return datafusion_common::not_impl_err!(
@@ -1084,12 +1115,69 @@ impl ListingTable {
         if !file_groups.is_empty() {
             file_groups.resize_with(file_group_count, || FileGroup::new(vec![]));
         }
-        let file_groups =
+        let mut file_groups =
             self.filter_declared_file_groups_by_partition_filters(file_groups, filters)?;
 
+        if matches!(output_partitioning, LogicalPartitioning::Range(_)) {
+            self.prune_range_file_groups(
+                ctx,
+                output_partitioning,
+                &mut file_groups,
+                range_filters,
+            )?;
+        }
+
         self.list_files_result_from_groups(ctx, file_groups, inexact_stats, false, vec![])
     }
 
+    fn prune_range_file_groups(
+        &self,
+        ctx: &dyn Session,
+        output_partitioning: &LogicalPartitioning,
+        file_groups: &mut [FileGroup],
+        filters: &[Expr],
+    ) -> datafusion_common::Result<()> {
+        // Do not invoke volatile expressions while planning. Unsupported
+        // expressions and pruning evaluation errors must keep the partitions.
+        let Some(filter) =
+            conjunction(filters.iter().filter(|f| !f.is_volatile()).cloned())
+        else {
+            return Ok(());
+        };
+        let schema = DFSchema::try_from(Arc::clone(&self.table_schema))?;
+        let planning_ctx = PhysicalPlanningContext::default();
+        let Partitioning::Range(range) = create_physical_partitioning(
+            output_partitioning,
+            &schema,
+            ctx.execution_props(),
+            &planning_ctx,
+        )?
+        else {
+            return Ok(());
+        };
+        let retained =
+            create_physical_expr(&filter, &schema, ctx.execution_props(), &planning_ctx)
+                .and_then(|expr| {
+                    PruningPredicateBuilder::new()
+                        .with_file_schema(Arc::clone(&self.table_schema))
+                        .try_build(expr)
+                })
+                .and_then(|predicate| {
+                    predicate.prune(&RangePruningStatistics::new(&range))
+                });
+        match retained {
+            Ok(retained) => {
+                for (group, keep) in file_groups.iter_mut().zip(retained) {
+                    if !keep {
+                        *group = FileGroup::new(vec![]);
+                    }
+                }
+            }
+            Err(error) => log::debug!("Skipping range partition pruning: {error}"),
+        }
+        Ok(())
+    }
+
     fn filter_declared_file_groups_by_partition_filters(
         &self,
         file_groups: Vec<FileGroup>,

datafusion/core/src/datasource/listing/table.rs

@@ -1431,6 +1431,75 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_list_files_prunes_normalized_range_boundaries() -> Result<()> {
+        let files = [
+            "bucket/key-prefix/file0",
+            "bucket/key-prefix/file1",
+            "bucket/key-prefix/file2",
+            "bucket/key-prefix/file3",
+        ];
+        let ctx = SessionContext::new();
+        // Logical split points need not have the physical range column's type.
+        let opt = ListingOptions::new(Arc::new(JsonFormat::default()))
+            .with_file_extension_opt(Some(""))
+            .with_output_partitioning(Some(LogicalPartitioning::Range(
+                LogicalRangePartitioning::try_new(
+                    vec![col("a").sort(true, true)],
+                    [10_i64, 20, 30]
+                        .into_iter()
+                        .map(|value| SplitPoint::new(vec![ScalarValue::from(value)]))
+                        .collect(),
+                )?,
+            )));
+        let table = listing_table_with_files(
+            &ctx,
+            &files,
+            "test:///bucket/key-prefix/",
+            opt,
+            Schema::new(vec![Field::new("a", DataType::Int32, true)]),
+        )?;
+
+        for (filter, expected) in [
+            (col("a").lt(lit(10_i32)), vec![1, 0, 0, 0]),
+            (col("a").eq(lit(10_i32)), vec![0, 1, 0, 0]),
+            (col("a").gt_eq(lit(30_i32)), vec![0, 0, 0, 1]),
+            (col("a").is_null(), vec![1, 0, 0, 0]),
+        ] {
+            let listed = table
+                .list_files_for_scan(&ctx.state(), std::slice::from_ref(&filter), Some(1))
+                .await?;
+            assert!(!listed.grouped_by_partition);
+            assert!(listed.partition_split_points.is_empty());
+            assert_eq!(
+                listed
+                    .file_groups
+                    .iter()
+                    .map(|g| g.len())
+                    .collect::<Vec<_>>(),
+                expected,
+                "{filter}"
+            );
+            let scan = table.scan(&ctx.state(), None, &[filter], None).await?;
+            assert_eq!(scan.output_partitioning().partition_count(), 4);
+        }
+        // Invalid pruning expressions are conservative; the normal scan/filter
+        // planning remains responsible for reporting any expression error.
+        let invalid_filter = col("missing").eq(lit(10_i32));
+        let listed = table
+            .list_files_for_scan(&ctx.state(), &[invalid_filter], None)
+            .await?;
+        assert_eq!(
+            listed
+                .file_groups
+                .iter()
+                .map(|g| g.len())
+                .collect::<Vec<_>>(),
+            vec![1, 1, 1, 1]
+        );
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_range_output_partitioning_normalizes_split_point_types() -> Result<()> {
         let files = ["bucket/key-prefix/file0"];

datafusion/pruning/src/lib.rs

@@ -22,10 +22,12 @@ mod in_list;
 mod key_range_bitmap_expr;
 mod primitive_in_list;
 mod pruning_predicate;
+mod range;
 mod string_in_list;
 
 pub use file_pruner::FilePruner;
 pub use pruning_predicate::{
     MAX_IN_LIST_SIZE, PredicateRewriter, PruningPredicate, PruningPredicateBuilder,
     PruningStatistics, RequiredColumns, UnhandledPredicateHook, build_pruning_predicate,
 };
+pub use range::RangePruningStatistics;