fix, add test

adriangb · adriangb · commit 0e03bdc46456 · 2025-05-12T01:08:37.000+02:00
diff --git a/datafusion/core/src/datasource/physical_plan/arrow_file.rs b/datafusion/core/src/datasource/physical_plan/arrow_file.rs
@@ -31,6 +31,7 @@ use datafusion_common::{Constraints, Statistics};
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
 use datafusion_datasource::source::DataSourceExec;
+use datafusion_datasource::PartitionedFile;
 use datafusion_datasource_json::source::JsonSource;
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
@@ -269,7 +270,11 @@ pub struct ArrowOpener {
 }
 
 impl FileOpener for ArrowOpener {
-    fn open(&self, file_meta: FileMeta, file: PartitionedFile) -> Result<FileOpenFuture> {
+    fn open(
+        &self,
+        file_meta: FileMeta,
+        _file: PartitionedFile,
+    ) -> Result<FileOpenFuture> {
         let object_store = Arc::clone(&self.object_store);
         let projection = self.projection.clone();
         Ok(Box::pin(async move {
diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs
@@ -67,7 +67,11 @@ pub struct TestOpener {
 }
 
 impl FileOpener for TestOpener {
-    fn open(&self, _file_meta: FileMeta, file: PartitionedFile) -> Result<FileOpenFuture> {
+    fn open(
+        &self,
+        _file_meta: FileMeta,
+        file: PartitionedFile,
+    ) -> Result<FileOpenFuture> {
         let mut batches = self.batches.clone();
         if let Some(batch_size) = self.batch_size {
             let batch = concat_batches(&batches[0].schema(), &batches)?;
diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs
@@ -129,7 +129,9 @@ mod private {
     use super::*;
 
     use bytes::Buf;
-    use datafusion_datasource::{file_meta::FileMeta, file_stream::FileOpenFuture, PartitionedFile};
+    use datafusion_datasource::{
+        file_meta::FileMeta, file_stream::FileOpenFuture, PartitionedFile,
+    };
     use futures::StreamExt;
     use object_store::{GetResultPayload, ObjectStore};
 
@@ -139,7 +141,11 @@ mod private {
     }
 
     impl FileOpener for AvroOpener {
-        fn open(&self, file_meta: FileMeta, _file: PartitionedFile) -> Result<FileOpenFuture> {
+        fn open(
+            &self,
+            file_meta: FileMeta,
+            _file: PartitionedFile,
+        ) -> Result<FileOpenFuture> {
             let config = Arc::clone(&self.config);
             let object_store = Arc::clone(&self.object_store);
             Ok(Box::pin(async move {
diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs
@@ -28,7 +28,7 @@ use datafusion_datasource::file_compression_type::FileCompressionType;
 use datafusion_datasource::file_meta::FileMeta;
 use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
 use datafusion_datasource::{
-    calculate_range, FileRange, ListingTableUrl, PartitionedFile, RangeCalculation
+    calculate_range, FileRange, ListingTableUrl, PartitionedFile, RangeCalculation,
 };
 
 use arrow::csv;
@@ -300,7 +300,11 @@ impl FileOpener for CsvOpener {
     ///  A,1,2,3,4,5,6,7,8,9\n
     ///  A},1,2,3,4,5,6,7,8,9\n
     ///  The lines read would be: [1, 2]
-    fn open(&self, file_meta: FileMeta, _file: PartitionedFile) -> Result<FileOpenFuture> {
+    fn open(
+        &self,
+        file_meta: FileMeta,
+        _file: PartitionedFile,
+    ) -> Result<FileOpenFuture> {
         // `self.config.has_header` controls whether to skip reading the 1st line header
         // If the .csv file is read in parallel and this `CsvOpener` is only reading some middle
         // partition, then don't skip first line
diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs
@@ -30,7 +30,9 @@ use datafusion_datasource::decoder::{deserialize_stream, DecoderDeserializer};
 use datafusion_datasource::file_compression_type::FileCompressionType;
 use datafusion_datasource::file_meta::FileMeta;
 use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
-use datafusion_datasource::{calculate_range, ListingTableUrl, PartitionedFile, RangeCalculation};
+use datafusion_datasource::{
+    calculate_range, ListingTableUrl, PartitionedFile, RangeCalculation,
+};
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 
 use arrow::json::ReaderBuilder;
@@ -328,7 +330,11 @@ impl FileOpener for JsonOpener {
     /// are applied to determine which lines to read:
     /// 1. The first line of the partition is the line in which the index of the first character >= `start`.
     /// 2. The last line of the partition is the line in which the byte at position `end - 1` resides.
-    fn open(&self, file_meta: FileMeta, _file: PartitionedFile) -> Result<FileOpenFuture> {
+    fn open(
+        &self,
+        file_meta: FileMeta,
+        _file: PartitionedFile,
+    ) -> Result<FileOpenFuture> {
         let store = Arc::clone(&self.object_store);
         let schema = Arc::clone(&self.projected_schema);
         let batch_size = self.batch_size;
diff --git a/datafusion/datasource-parquet/src/opener.rs b/datafusion/datasource-parquet/src/opener.rs
@@ -92,24 +92,35 @@ impl FileOpener for ParquetOpener {
         // Prune using known statistics
         match (&file.statistics, &self.predicate) {
             (Some(stats), Some(predicate)) => {
-                let pruning_predicate =
-                    build_pruning_predicate(Arc::clone(predicate), &self.table_schema, &predicate_creation_errors);
+                let pruning_predicate = build_pruning_predicate(
+                    Arc::clone(predicate),
+                    &self.table_schema,
+                    &predicate_creation_errors,
+                );
                 if let Some(pruning_predicate) = pruning_predicate {
                     let pruning_stats = PrunableStatistics::new(
-                        Arc::clone(stats),
+                        vec![Arc::clone(stats)],
                         Arc::clone(&self.table_schema),
                     );
                     match pruning_predicate.prune(&pruning_stats) {
                         Ok(values) => {
                             // We expect a single container -> if all containers are false skip this file
                             if values.into_iter().all(|v| !v) {
                                 // Return an empty stream
-                                todo!()
+                                return Ok(Box::pin(async move {
+                                    Ok(futures::stream::empty().boxed())
+                                }));
                             }
                         }
                         // stats filter array could not be built, so we can't prune
                         Err(e) => {
-                            log::debug!("Error evaluating row group predicate values {e}");
+                            let err = format!(
+                                "Error evaluating row group predicate values {e}"
+                            );
+                            println!("{err}");
+                            log::debug!(
+                                "Error evaluating row group predicate values {e}"
+                            );
                             predicate_creation_errors.add(1);
                         }
                     }
@@ -469,3 +480,145 @@ async fn load_page_index<T: AsyncFileReader>(
         Ok(reader_metadata)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use std::sync::Arc;
+
+    use bytes::{BufMut, BytesMut};
+    use chrono::Utc;
+    use datafusion_common::{
+        record_batch, stats::Precision, ColumnStatistics, ScalarValue, Statistics,
+    };
+    use datafusion_datasource::{
+        file_meta::FileMeta, file_stream::FileOpener,
+        schema_adapter::DefaultSchemaAdapterFactory, PartitionedFile,
+    };
+    use datafusion_expr::{col, lit};
+    use datafusion_physical_expr::planner::logical2physical;
+    use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
+    use futures::{Stream, StreamExt};
+    use object_store::{memory::InMemory, path::Path, ObjectMeta, ObjectStore};
+    use parquet::arrow::ArrowWriter;
+
+    use crate::{opener::ParquetOpener, DefaultParquetFileReaderFactory};
+
+    async fn count_batches_and_rows(
+        mut stream: std::pin::Pin<
+            Box<
+                dyn Stream<
+                        Item = Result<
+                            arrow::array::RecordBatch,
+                            arrow::error::ArrowError,
+                        >,
+                    > + Send,
+            >,
+        >,
+    ) -> (usize, usize) {
+        let mut num_batches = 0;
+        let mut num_rows = 0;
+        while let Some(Ok(batch)) = stream.next().await {
+            num_rows += batch.num_rows();
+            num_batches += 1;
+        }
+        (num_batches, num_rows)
+    }
+
+    #[tokio::test]
+    async fn test_prune_based_on_statistics() {
+        let batch = record_batch!(
+            ("a", Int32, vec![Some(1), Some(2), Some(2)]),
+            ("b", Float32, vec![Some(1.0), Some(2.0), None])
+        )
+        .unwrap();
+
+        let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+        let mut out = BytesMut::new().writer();
+        {
+            let mut writer =
+                ArrowWriter::try_new(&mut out, batch.schema(), None).unwrap();
+            writer.write(&batch).unwrap();
+            writer.finish().unwrap();
+        }
+        let data = out.into_inner().freeze();
+        let data_size = data.len();
+        store
+            .put(&Path::from("test.parquet"), data.into())
+            .await
+            .unwrap();
+
+        let schema = batch.schema();
+        let file = PartitionedFile::new(
+            "file.parquet".to_string(),
+            u64::try_from(data_size).unwrap(),
+        )
+        .with_statistics(Arc::new(
+            Statistics::new_unknown(&schema)
+                .add_column_statistics(ColumnStatistics::new_unknown())
+                .add_column_statistics(
+                    ColumnStatistics::new_unknown()
+                        .with_min_value(Precision::Exact(ScalarValue::Float32(Some(1.0))))
+                        .with_max_value(Precision::Exact(ScalarValue::Float32(Some(2.0))))
+                        .with_null_count(Precision::Exact(1)),
+                ),
+        ));
+
+        let make_opener = |predicate| {
+            ParquetOpener {
+                partition_index: 0,
+                projection: Arc::new([0, 1]),
+                batch_size: 1024,
+                limit: None,
+                predicate: Some(predicate),
+                table_schema: schema.clone(),
+                metadata_size_hint: None,
+                metrics: ExecutionPlanMetricsSet::new(),
+                parquet_file_reader_factory: Arc::new(
+                    DefaultParquetFileReaderFactory::new(Arc::clone(&store)),
+                ),
+                pushdown_filters: false, // note that this is false!
+                reorder_filters: false,
+                enable_page_index: false,
+                enable_bloom_filter: false,
+                schema_adapter_factory: Arc::new(DefaultSchemaAdapterFactory),
+                enable_row_group_stats_pruning: true,
+                coerce_int96: None,
+            }
+        };
+
+        let make_meta = || FileMeta {
+            object_meta: ObjectMeta {
+                location: Path::from("test.parquet"),
+                last_modified: Utc::now(),
+                size: u64::try_from(data_size).unwrap(),
+                e_tag: None,
+                version: None,
+            },
+            range: None,
+            extensions: None,
+            metadata_size_hint: None,
+        };
+
+        // A filter on "a" should not exclude any rows even if it matches the data
+        let expr = col("a").eq(lit(1));
+        let predicate = logical2physical(&expr, &schema);
+        let opener = make_opener(predicate);
+        let stream = opener
+            .open(make_meta(), file.clone())
+            .unwrap()
+            .await
+            .unwrap();
+        let (num_batches, num_rows) = count_batches_and_rows(stream).await;
+        assert_eq!(num_batches, 1);
+        assert_eq!(num_rows, 3);
+
+        // A filter on `b = 5.0` should exclude all rows
+        let expr = col("b").eq(lit(ScalarValue::Float32(Some(5.0))));
+        let predicate = logical2physical(&expr, &schema);
+        let opener = make_opener(predicate);
+        let stream = opener.open(make_meta(), file).unwrap().await.unwrap();
+        let (num_batches, num_rows) = count_batches_and_rows(stream).await;
+        assert_eq!(num_batches, 0);
+        assert_eq!(num_rows, 0);
+    }
+}
diff --git a/datafusion/datasource/src/file_stream.rs b/datafusion/datasource/src/file_stream.rs
@@ -556,7 +556,11 @@ mod tests {
     }
 
     impl FileOpener for TestOpener {
-        fn open(&self, _file_meta: FileMeta, _file: PartitionedFile) -> Result<FileOpenFuture> {
+        fn open(
+            &self,
+            _file_meta: FileMeta,
+            _file: PartitionedFile,
+        ) -> Result<FileOpenFuture> {
             let idx = self.current_idx.fetch_add(1, Ordering::SeqCst);
 
             if self.error_opening_idx.contains(&idx) {
diff --git a/datafusion/physical-optimizer/src/pruning.rs b/datafusion/physical-optimizer/src/pruning.rs
