Implement Parquet filter pushdown via new filter pushdown APIs

Author: adriangb

datafusion-examples/examples/advanced_parquet_index.rs

@@ -495,7 +495,7 @@ impl TableProvider for IndexTableProvider {
             ParquetSource::default()
                 // provide the predicate so the DataSourceExec can try and prune
                 // row groups internally
-                .with_predicate(Arc::clone(&schema), predicate)
+                .with_predicate(predicate)
                 // provide the factory to create parquet reader without re-reading metadata
                 .with_parquet_file_reader_factory(Arc::new(reader_factory)),
         );

datafusion-examples/examples/custom_file_format.rs

@@ -22,7 +22,6 @@ use arrow::{
     datatypes::{DataType, Field, Schema, SchemaRef, UInt64Type},
 };
 use datafusion::physical_expr::LexRequirement;
-use datafusion::physical_expr::PhysicalExpr;
 use datafusion::{
     catalog::Session,
     common::{GetExt, Statistics},
@@ -112,11 +111,8 @@ impl FileFormat for TSVFileFormat {
         &self,
         state: &dyn Session,
         conf: FileScanConfig,
-        filters: Option<&Arc<dyn PhysicalExpr>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        self.csv_file_format
-            .create_physical_plan(state, conf, filters)
-            .await
+        self.csv_file_format.create_physical_plan(state, conf).await
     }
 
     async fn create_writer_physical_plan(

datafusion-examples/examples/parquet_index.rs

@@ -242,8 +242,7 @@ impl TableProvider for IndexTableProvider {
         let files = self.index.get_files(predicate.clone())?;
 
         let object_store_url = ObjectStoreUrl::parse("file://")?;
-        let source =
-            Arc::new(ParquetSource::default().with_predicate(self.schema(), predicate));
+        let source = Arc::new(ParquetSource::default().with_predicate(predicate));
         let mut file_scan_config_builder =
             FileScanConfigBuilder::new(object_store_url, self.schema(), source)
                 .with_projection(projection.cloned())

datafusion/core/src/datasource/file_format/arrow.rs

@@ -54,7 +54,6 @@ use datafusion_datasource::sink::{DataSink, DataSinkExec};
 use datafusion_datasource::write::ObjectWriterBuilder;
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_expr::dml::InsertOp;
-use datafusion_physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::sort_expr::LexRequirement;
 
 use async_trait::async_trait;
@@ -174,7 +173,6 @@ impl FileFormat for ArrowFormat {
         &self,
         _state: &dyn Session,
         conf: FileScanConfig,
-        _filters: Option<&Arc<dyn PhysicalExpr>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let source = Arc::new(ArrowSource::default());
         let config = FileScanConfigBuilder::from(conf)

datafusion/core/src/datasource/file_format/mod.rs

@@ -93,7 +93,6 @@ pub(crate) mod test_util {
                 .with_projection(projection)
                 .with_limit(limit)
                 .build(),
-                None,
             )
             .await?;
         Ok(exec)

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

@@ -24,9 +24,7 @@ use std::{any::Any, str::FromStr, sync::Arc};
 
 use crate::datasource::{
     create_ordering,
-    file_format::{
-        file_compression_type::FileCompressionType, FileFormat, FilePushdownSupport,
-    },
+    file_format::{file_compression_type::FileCompressionType, FileFormat},
     physical_plan::FileSinkConfig,
 };
 use crate::execution::context::SessionState;
@@ -35,22 +33,19 @@ use datafusion_common::{config_err, DataFusionError, Result};
 use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
 use datafusion_datasource::schema_adapter::DefaultSchemaAdapterFactory;
 use datafusion_expr::dml::InsertOp;
-use datafusion_expr::{utils::conjunction, Expr, TableProviderFilterPushDown};
+use datafusion_expr::{Expr, TableProviderFilterPushDown};
 use datafusion_expr::{SortExpr, TableType};
 use datafusion_physical_plan::empty::EmptyExec;
 use datafusion_physical_plan::{ExecutionPlan, Statistics};
 
 use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, SchemaRef};
 use datafusion_common::{
-    config_datafusion_err, internal_err, plan_err, project_schema, Constraints,
-    SchemaExt, ToDFSchema,
+    config_datafusion_err, internal_err, plan_err, project_schema, Constraints, SchemaExt,
 };
 use datafusion_execution::cache::{
     cache_manager::FileStatisticsCache, cache_unit::DefaultFileStatisticsCache,
 };
-use datafusion_physical_expr::{
-    create_physical_expr, LexOrdering, PhysicalSortRequirement,
-};
+use datafusion_physical_expr::{LexOrdering, PhysicalSortRequirement};
 
 use async_trait::async_trait;
 use datafusion_catalog::Session;
@@ -941,19 +936,6 @@ impl TableProvider for ListingTable {
             None => {} // no ordering required
         };
 
-        let filters = match conjunction(filters.to_vec()) {
-            Some(expr) => {
-                let table_df_schema = self.table_schema.as_ref().clone().to_dfschema()?;
-                let filters = create_physical_expr(
-                    &expr,
-                    &table_df_schema,
-                    state.execution_props(),
-                )?;
-                Some(filters)
-            }
-            None => None,
-        };
-
         let Some(object_store_url) =
             self.table_paths.first().map(ListingTableUrl::object_store)
         else {
@@ -978,7 +960,6 @@ impl TableProvider for ListingTable {
                 .with_output_ordering(output_ordering)
                 .with_table_partition_cols(table_partition_cols)
                 .build(),
-                filters.as_ref(),
             )
             .await
     }
@@ -1002,18 +983,6 @@ impl TableProvider for ListingTable {
                     return Ok(TableProviderFilterPushDown::Exact);
                 }
 
-                // if we can't push it down completely with only the filename-based/path-based
-                // column names, then we should check if we can do parquet predicate pushdown
-                let supports_pushdown = self.options.format.supports_filters_pushdown(
-                    &self.file_schema,
-                    &self.table_schema,
-                    &[filter],
-                )?;
-
-                if supports_pushdown == FilePushdownSupport::Supported {
-                    return Ok(TableProviderFilterPushDown::Exact);
-                }
-
                 Ok(TableProviderFilterPushDown::Inexact)
             })
             .collect()

datafusion/core/src/datasource/physical_plan/parquet.rs

@@ -54,6 +54,7 @@ mod tests {
     use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
     use datafusion_datasource::source::DataSourceExec;
 
+    use datafusion_datasource::file::FileSource;
     use datafusion_datasource::{FileRange, PartitionedFile};
     use datafusion_datasource_parquet::source::ParquetSource;
     use datafusion_datasource_parquet::{
@@ -139,7 +140,7 @@ mod tests {
             self.round_trip(batches).await.batches
         }
 
-        fn build_file_source(&self, file_schema: SchemaRef) -> Arc<ParquetSource> {
+        fn build_file_source(&self, file_schema: SchemaRef) -> Arc<dyn FileSource> {
             // set up predicate (this is normally done by a layer higher up)
             let predicate = self
                 .predicate
@@ -148,7 +149,7 @@ mod tests {
 
             let mut source = ParquetSource::default();
             if let Some(predicate) = predicate {
-                source = source.with_predicate(Arc::clone(&file_schema), predicate);
+                source = source.with_predicate(predicate);
             }
 
             if self.pushdown_predicate {
@@ -161,14 +162,14 @@ mod tests {
                 source = source.with_enable_page_index(true);
             }
 
-            Arc::new(source)
+            source.with_schema(Arc::clone(&file_schema))
         }
 
         fn build_parquet_exec(
             &self,
             file_schema: SchemaRef,
             file_group: FileGroup,
-            source: Arc<ParquetSource>,
+            source: Arc<dyn FileSource>,
         ) -> Arc<DataSourceExec> {
             let base_config = FileScanConfigBuilder::new(
                 ObjectStoreUrl::local_filesystem(),

datafusion/core/src/test_util/parquet.rs

@@ -37,6 +37,7 @@ use crate::physical_plan::metrics::MetricsSet;
 use crate::physical_plan::ExecutionPlan;
 use crate::prelude::{Expr, SessionConfig, SessionContext};
 
+use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
 use datafusion_datasource::source::DataSourceExec;
 use object_store::path::Path;
@@ -182,10 +183,11 @@ impl TestParquetFile {
             let physical_filter_expr =
                 create_physical_expr(&filter, &df_schema, &ExecutionProps::default())?;
 
-            let source = Arc::new(ParquetSource::new(parquet_options).with_predicate(
-                Arc::clone(&self.schema),
-                Arc::clone(&physical_filter_expr),
-            ));
+            let source = Arc::new(
+                ParquetSource::new(parquet_options)
+                    .with_predicate(Arc::clone(&physical_filter_expr)),
+            )
+            .with_schema(Arc::clone(&self.schema));
             let config = scan_config_builder.with_source(source).build();
             let parquet_exec = DataSourceExec::from_data_source(config);
 

datafusion/core/tests/fuzz_cases/pruning.rs

@@ -276,7 +276,7 @@ async fn execute_with_predicate(
     ctx: &SessionContext,
 ) -> Vec<String> {
     let parquet_source = if prune_stats {
-        ParquetSource::default().with_predicate(Arc::clone(&schema), predicate.clone())
+        ParquetSource::default().with_predicate(predicate.clone())
     } else {
         ParquetSource::default()
     };

datafusion/core/tests/parquet/external_access_plan.rs

@@ -346,7 +346,7 @@ impl TestFull {
         let source = if let Some(predicate) = predicate {
             let df_schema = DFSchema::try_from(schema.clone())?;
             let predicate = ctx.create_physical_expr(predicate, &df_schema)?;
-            Arc::new(ParquetSource::default().with_predicate(schema.clone(), predicate))
+            Arc::new(ParquetSource::default().with_predicate(predicate))
         } else {
             Arc::new(ParquetSource::default())
         };

datafusion/core/tests/parquet/file_statistics.rs

@@ -28,6 +28,7 @@ use datafusion::execution::context::SessionState;
 use datafusion::execution::session_state::SessionStateBuilder;
 use datafusion::prelude::SessionContext;
 use datafusion_common::stats::Precision;
+use datafusion_common::DFSchema;
 use datafusion_execution::cache::cache_manager::CacheManagerConfig;
 use datafusion_execution::cache::cache_unit::{
     DefaultFileStatisticsCache, DefaultListFilesCache,
@@ -37,6 +38,10 @@ use datafusion_execution::runtime_env::RuntimeEnvBuilder;
 use datafusion_expr::{col, lit, Expr};
 
 use datafusion::datasource::physical_plan::FileScanConfig;
+use datafusion_physical_optimizer::filter_pushdown::FilterPushdown;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
+use datafusion_physical_plan::filter::FilterExec;
+use datafusion_physical_plan::ExecutionPlan;
 use tempfile::tempdir;
 
 #[tokio::test]
@@ -47,21 +52,49 @@ async fn check_stats_precision_with_filter_pushdown() {
 
     let opt = ListingOptions::new(Arc::new(ParquetFormat::default()));
     let table = get_listing_table(&table_path, None, &opt).await;
+
     let (_, _, state) = get_cache_runtime_state();
+    let mut options = state.config().options().clone();
+    options.execution.parquet.pushdown_filters = true;
+
     // Scan without filter, stats are exact
     let exec = table.scan(&state, None, &[], None).await.unwrap();
     assert_eq!(
         exec.partition_statistics(None).unwrap().num_rows,
-        Precision::Exact(8)
+        Precision::Exact(8),
+        "Stats without filter should be exact"
     );
 
-    // Scan with filter pushdown, stats are inexact
-    let filter = Expr::gt(col("id"), lit(1));
+    // This is a filter that cannot be evaluated by the table provider scanning
+    // (it is not a partition filter). Therefore; it will be pushed down to the
+    // source operator after the appropriate optimizer pass.
+    let filter_expr = Expr::gt(col("id"), lit(1));
+    let exec_with_filter = table
+        .scan(&state, None, &[filter_expr.clone()], None)
+        .await
+        .unwrap();
+
+    let ctx = SessionContext::new();
+    let df_schema = DFSchema::try_from(table.schema()).unwrap();
+    let physical_filter = ctx.create_physical_expr(filter_expr, &df_schema).unwrap();
 
-    let exec = table.scan(&state, None, &[filter], None).await.unwrap();
+    let filtered_exec =
+        Arc::new(FilterExec::try_new(physical_filter, exec_with_filter).unwrap())
+            as Arc<dyn ExecutionPlan>;
+
+    let optimized_exec = FilterPushdown::new()
+        .optimize(filtered_exec, &options)
+        .unwrap();
+
+    assert!(
+        optimized_exec.as_any().is::<DataSourceExec>(),
+        "Sanity check that the pushdown did what we expected"
+    );
+    // Scan with filter pushdown, stats are inexact
     assert_eq!(
-        exec.partition_statistics(None).unwrap().num_rows,
-        Precision::Inexact(8)
+        optimized_exec.partition_statistics(None).unwrap().num_rows,
+        Precision::Inexact(8),
+        "Stats after filter pushdown should be inexact"
     );
 }
 

datafusion/core/tests/parquet/page_pruning.rs

@@ -77,7 +77,7 @@ async fn get_parquet_exec(state: &SessionState, filter: Expr) -> DataSourceExec
 
     let source = Arc::new(
         ParquetSource::default()
-            .with_predicate(Arc::clone(&schema), predicate)
+            .with_predicate(predicate)
             .with_enable_page_index(true),
     );
     let base_config = FileScanConfigBuilder::new(object_store_url, schema, source)

datafusion/core/tests/sql/path_partition.rs

@@ -25,8 +25,6 @@ use std::sync::Arc;
 
 use arrow::datatypes::DataType;
 use datafusion::datasource::listing::ListingTableUrl;
-use datafusion::datasource::physical_plan::ParquetSource;
-use datafusion::datasource::source::DataSourceExec;
 use datafusion::{
     datasource::{
         file_format::{csv::CsvFormat, parquet::ParquetFormat},
@@ -42,8 +40,6 @@ use datafusion_common::stats::Precision;
 use datafusion_common::test_util::batches_to_sort_string;
 use datafusion_common::ScalarValue;
 use datafusion_execution::config::SessionConfig;
-use datafusion_expr::{col, lit, Expr, Operator};
-use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
 
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -57,55 +53,6 @@ use object_store::{
 use object_store::{Attributes, MultipartUpload, PutMultipartOpts, PutPayload};
 use url::Url;
 
-#[tokio::test]
-async fn parquet_partition_pruning_filter() -> Result<()> {
-    let ctx = SessionContext::new();
-
-    let table = create_partitioned_alltypes_parquet_table(
-        &ctx,
-        &[
-            "year=2021/month=09/day=09/file.parquet",
-            "year=2021/month=10/day=09/file.parquet",
-            "year=2021/month=10/day=28/file.parquet",
-        ],
-        &[
-            ("year", DataType::Int32),
-            ("month", DataType::Int32),
-            ("day", DataType::Int32),
-        ],
-        "mirror:///",
-        "alltypes_plain.parquet",
-    )
-    .await;
-
-    // The first three filters can be resolved using only the partition columns.
-    let filters = [
-        Expr::eq(col("year"), lit(2021)),
-        Expr::eq(col("month"), lit(10)),
-        Expr::eq(col("day"), lit(28)),
-        Expr::gt(col("id"), lit(1)),
-    ];
-    let exec = table.scan(&ctx.state(), None, &filters, None).await?;
-    let data_source_exec = exec.as_any().downcast_ref::<DataSourceExec>().unwrap();
-    if let Some((_, parquet_config)) =
-        data_source_exec.downcast_to_file_source::<ParquetSource>()
-    {
-        let pred = parquet_config.predicate().unwrap();
-        // Only the last filter should be pushdown to TableScan
-        let expected = Arc::new(BinaryExpr::new(
-            Arc::new(Column::new_with_schema("id", &exec.schema()).unwrap()),
-            Operator::Gt,
-            Arc::new(Literal::new(ScalarValue::Int32(Some(1)))),
-        ));
-
-        assert!(pred.as_any().is::<BinaryExpr>());
-        let pred = pred.as_any().downcast_ref::<BinaryExpr>().unwrap();
-
-        assert_eq!(pred, expected.as_ref());
-    }
-    Ok(())
-}
-
 #[tokio::test]
 async fn parquet_distinct_partition_col() -> Result<()> {
     let ctx = SessionContext::new();