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/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/benches/push_down_filter.rs

@@ -0,0 +1,125 @@
+// 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.
+
+use arrow::array::RecordBatch;
+use arrow::datatypes::{DataType, Field, Schema};
+use bytes::{BufMut, BytesMut};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use datafusion::config::ConfigOptions;
+use datafusion::prelude::{ParquetReadOptions, SessionContext};
+use datafusion_execution::object_store::ObjectStoreUrl;
+use datafusion_physical_optimizer::push_down_filter::PushdownFilter;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
+use datafusion_physical_plan::ExecutionPlan;
+use object_store::memory::InMemory;
+use object_store::path::Path;
+use object_store::ObjectStore;
+use parquet::arrow::ArrowWriter;
+use std::sync::Arc;
+
+async fn create_plan() -> Arc<dyn ExecutionPlan> {
+    let ctx = SessionContext::new();
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Int32, true),
+        Field::new("name", DataType::Utf8, true),
+        Field::new("age", DataType::UInt16, true),
+        Field::new("salary", DataType::Float64, true),
+    ]));
+    let batch = RecordBatch::new_empty(schema);
+
+    let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+    let data = {
+        let out = BytesMut::new();
+        let mut writer =
+            ArrowWriter::try_new(out.writer(), batch.schema(), None).unwrap();
+        writer.write(&batch).unwrap();
+        writer.finish().unwrap();
+        writer.into_inner().unwrap().into_inner().freeze()
+    };
+    store
+        .put(&Path::from("test.parquet"), data.into())
+        .await
+        .unwrap();
+    ctx.register_object_store(
+        ObjectStoreUrl::parse("memory://").unwrap().as_ref(),
+        store,
+    );
+
+    ctx.register_parquet("t", "memory://", ParquetReadOptions::default())
+        .await
+        .unwrap();
+
+    let df = ctx
+        .sql(
+            r"
+        WITH brackets AS (
+            SELECT age % 10 AS age_bracket
+            FROM t
+            GROUP BY age % 10
+            HAVING COUNT(*) > 10
+        )
+        SELECT id, name, age, salary
+        FROM t
+        JOIN brackets ON t.age % 10 = brackets.age_bracket
+        WHERE age > 20 AND data.salary > 1000
+        ORDER BY data.salary DESC
+        LIMIT 100
+    ",
+        )
+        .await
+        .unwrap();
+
+    df.create_physical_plan().await.unwrap()
+}
+
+#[derive(Clone)]
+struct BenchmarkPlan {
+    plan: Arc<dyn ExecutionPlan>,
+    config: ConfigOptions,
+}
+
+impl std::fmt::Display for BenchmarkPlan {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "BenchmarkPlan")
+    }
+}
+
+fn bench_push_down_filter(c: &mut Criterion) {
+    // Create a relatively complex plan
+    let plan = tokio::runtime::Runtime::new()
+        .unwrap()
+        .block_on(create_plan());
+    let mut config = ConfigOptions::default();
+    config.execution.parquet.pushdown_filters = true;
+    let plan = BenchmarkPlan { plan, config };
+
+    c.bench_with_input(
+        BenchmarkId::new("push_down_filter", plan.clone()),
+        &plan,
+        |b, plan| {
+            b.iter(|| {
+                let optimizer = PushdownFilter::new();
+                optimizer
+                    .optimize(Arc::clone(&plan.plan), &plan.config)
+                    .unwrap();
+            });
+        },
+    );
+}
+
+criterion_group!(benches, bench_push_down_filter);
+criterion_main!(benches);

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;
@@ -1002,18 +1000,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/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/physical_optimizer/push_down_filter.rs

@@ -99,7 +99,7 @@ impl FileSource for TestSource {
     }
 
     fn with_schema(&self, _schema: SchemaRef) -> Arc<dyn FileSource> {
-        todo!("should not be called")
+        Arc::new(self.clone()) as Arc<dyn FileSource>
     }
 
     fn with_projection(&self, _config: &FileScanConfig) -> Arc<dyn FileSource> {

datafusion/datasource-parquet/src/file_format.rs

@@ -31,9 +31,7 @@ use datafusion_datasource::write::{
     get_writer_schema, ObjectWriterBuilder, SharedBuffer,
 };
 
-use datafusion_datasource::file_format::{
-    FileFormat, FileFormatFactory, FilePushdownSupport,
-};
+use datafusion_datasource::file_format::{FileFormat, FileFormatFactory};
 use datafusion_datasource::write::demux::DemuxedStreamReceiver;
 
 use arrow::compute::sum;
@@ -54,15 +52,13 @@ use datafusion_datasource::sink::{DataSink, DataSinkExec};
 use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_expr::dml::InsertOp;
-use datafusion_expr::Expr;
 use datafusion_functions_aggregate::min_max::{MaxAccumulator, MinAccumulator};
 use datafusion_physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::sort_expr::LexRequirement;
 use datafusion_physical_plan::Accumulator;
 use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan};
 use datafusion_session::Session;
 
-use crate::can_expr_be_pushed_down_with_schemas;
 use crate::source::{parse_coerce_int96_string, ParquetSource};
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -433,7 +429,7 @@ impl FileFormat for ParquetFormat {
         let mut source = ParquetSource::new(self.options.clone());
 
         if let Some(predicate) = predicate {
-            source = source.with_predicate(Arc::clone(&conf.file_schema), predicate);
+            source = source.with_predicate(predicate);
         }
         if let Some(metadata_size_hint) = metadata_size_hint {
             source = source.with_metadata_size_hint(metadata_size_hint)
@@ -461,27 +457,6 @@ impl FileFormat for ParquetFormat {
         Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
     }
 
-    fn supports_filters_pushdown(
-        &self,
-        file_schema: &Schema,
-        _table_schema: &Schema,
-        filters: &[&Expr],
-    ) -> Result<FilePushdownSupport> {
-        if !self.options().global.pushdown_filters {
-            return Ok(FilePushdownSupport::NoSupport);
-        }
-
-        let all_supported = filters
-            .iter()
-            .all(|filter| can_expr_be_pushed_down_with_schemas(filter, file_schema));
-
-        Ok(if all_supported {
-            FilePushdownSupport::Supported
-        } else {
-            FilePushdownSupport::NotSupportedForFilter
-        })
-    }
-
     fn file_source(&self) -> Arc<dyn FileSource> {
         Arc::new(ParquetSource::default())
     }

datafusion/datasource-parquet/src/mod.rs

@@ -59,7 +59,6 @@ pub use metrics::ParquetFileMetrics;
 pub use page_filter::PagePruningAccessPlanFilter;
 pub use reader::{DefaultParquetFileReaderFactory, ParquetFileReaderFactory};
 pub use row_filter::build_row_filter;
-pub use row_filter::can_expr_be_pushed_down_with_schemas;
 pub use row_group_filter::RowGroupAccessPlanFilter;
 use source::ParquetSource;
 pub use writer::plan_to_parquet;
@@ -223,8 +222,7 @@ impl ParquetExecBuilder {
         } = self;
         let mut parquet = ParquetSource::new(table_parquet_options);
         if let Some(predicate) = predicate.clone() {
-            parquet = parquet
-                .with_predicate(Arc::clone(&file_scan_config.file_schema), predicate);
+            parquet = parquet.with_predicate(predicate);
         }
         if let Some(metadata_size_hint) = metadata_size_hint {
             parquet = parquet.with_metadata_size_hint(metadata_size_hint)
@@ -244,7 +242,7 @@ impl ParquetExecBuilder {
             inner: DataSourceExec::new(Arc::new(base_config.clone())),
             base_config,
             predicate,
-            pruning_predicate: parquet.pruning_predicate,
+            pruning_predicate: None, // for backwards compat since `ParquetExec` is only for backwards compat anyway
             schema_adapter_factory: parquet.schema_adapter_factory,
             parquet_file_reader_factory: parquet.parquet_file_reader_factory,
             table_parquet_options: parquet.table_parquet_options,

datafusion/datasource-parquet/src/opener.rs

@@ -178,7 +178,7 @@ impl FileOpener for ParquetOpener {
 
             // Build predicates for this specific file
             let (pruning_predicate, page_pruning_predicate) = build_pruning_predicates(
-                &predicate,
+                predicate.as_ref(),
                 &physical_file_schema,
                 &predicate_creation_errors,
             );
@@ -390,8 +390,8 @@ pub(crate) fn build_page_pruning_predicate(
     ))
 }
 
-fn build_pruning_predicates(
-    predicate: &Option<Arc<dyn PhysicalExpr>>,
+pub(crate) fn build_pruning_predicates(
+    predicate: Option<&Arc<dyn PhysicalExpr>>,
     file_schema: &SchemaRef,
     predicate_creation_errors: &Count,
 ) -> (