Thread through max_predicate_cache_size, add test

Author: alamb

datafusion/common/src/config.rs

@@ -565,6 +565,14 @@ config_namespace! {
         /// (reading) Use any available bloom filters when reading parquet files
         pub bloom_filter_on_read: bool, default = true
 
+        /// (reading) The maximum predicate cache size, in bytes. When
+        /// `pushdown_filters` is enabled, sets the maximum memory used to cache
+        /// the results of predicate evaluation between filter evaluation and
+        /// output generation. Decreasing this value will reduce memory usage,
+        /// but may increase IO and CPU usage. None means use the default
+        /// parquet reader setting. 0 means no caching.
+        pub max_predicate_cache_size: Option<usize>, default = None
+
         // The following options affect writing to parquet files
         // and map to parquet::file::properties::WriterProperties
 

datafusion/common/src/file_options/parquet_writer.rs

@@ -233,6 +233,7 @@ impl ParquetOptions {
             binary_as_string: _, // not used for writer props
             coerce_int96: _,     // not used for writer props
             skip_arrow_metadata: _,
+            max_predicate_cache_size: _,
         } = self;
 
         let mut builder = WriterProperties::builder()
@@ -425,6 +426,10 @@ pub(crate) fn parse_statistics_string(str_setting: &str) -> Result<EnabledStatis
 #[cfg(feature = "parquet")]
 #[cfg(test)]
 mod tests {
+    use super::*;
+    use crate::config::{ParquetColumnOptions, ParquetEncryptionOptions, ParquetOptions};
+    #[cfg(feature = "parquet_encryption")]
+    use crate::encryption::map_encryption_to_config_encryption;
     use parquet::{
         basic::Compression,
         file::properties::{
@@ -434,11 +439,6 @@ mod tests {
     };
     use std::collections::HashMap;
 
-    use super::*;
-    use crate::config::{ParquetColumnOptions, ParquetEncryptionOptions, ParquetOptions};
-    #[cfg(feature = "parquet_encryption")]
-    use crate::encryption::map_encryption_to_config_encryption;
-
     const COL_NAME: &str = "configured";
 
     /// Take the column defaults provided in [`ParquetOptions`], and generate a non-default col config.
@@ -500,6 +500,7 @@ mod tests {
             binary_as_string: defaults.binary_as_string,
             skip_arrow_metadata: defaults.skip_arrow_metadata,
             coerce_int96: None,
+            max_predicate_cache_size: defaults.max_predicate_cache_size,
         }
     }
 
@@ -606,6 +607,8 @@ mod tests {
                 maximum_buffered_record_batches_per_stream: global_options_defaults
                     .maximum_buffered_record_batches_per_stream,
                 bloom_filter_on_read: global_options_defaults.bloom_filter_on_read,
+                max_predicate_cache_size: global_options_defaults
+                    .max_predicate_cache_size,
                 schema_force_view_types: global_options_defaults.schema_force_view_types,
                 binary_as_string: global_options_defaults.binary_as_string,
                 skip_arrow_metadata: global_options_defaults.skip_arrow_metadata,

datafusion/core/tests/parquet/filter_pushdown.rs

@@ -26,8 +26,6 @@
 //! select * from data limit 10;
 //! ```
 
-use std::path::Path;
-
 use arrow::compute::concat_batches;
 use arrow::record_batch::RecordBatch;
 use datafusion::physical_plan::collect;
@@ -37,7 +35,10 @@ use datafusion::prelude::{
 };
 use datafusion::test_util::parquet::{ParquetScanOptions, TestParquetFile};
 use datafusion_expr::utils::{conjunction, disjunction, split_conjunction};
+use std::path::Path;
 
+use datafusion_common::test_util::parquet_test_data;
+use datafusion_execution::config::SessionConfig;
 use itertools::Itertools;
 use parquet::file::properties::WriterProperties;
 use tempfile::TempDir;
@@ -601,3 +602,99 @@ fn get_value(metrics: &MetricsSet, metric_name: &str) -> usize {
         }
     }
 }
+
+#[tokio::test]
+async fn predicate_cache_default() -> datafusion_common::Result<()> {
+    let ctx = SessionContext::new();
+    // The cache is on by default, but not used unless filter pushdown is enabled
+    PredicateCacheTest {
+        expected_inner_records: 0,
+        expected_records: 0,
+    }
+    .run(&ctx)
+    .await
+}
+
+#[tokio::test]
+async fn predicate_cache_pushdown_default() -> datafusion_common::Result<()> {
+    let mut config = SessionConfig::new();
+    config.options_mut().execution.parquet.pushdown_filters = true;
+    let ctx = SessionContext::new_with_config(config);
+    // The cache is on by default, and used when filter pushdown is enabled
+    PredicateCacheTest {
+        expected_inner_records: 8,
+        expected_records: 4,
+    }
+    .run(&ctx)
+    .await
+}
+
+#[tokio::test]
+async fn predicate_cache_pushdown_disable() -> datafusion_common::Result<()> {
+    // Can disable the cache even with filter pushdown by setting the size to 0. In this case we
+    // expect the inner records are reported but no records are read from the cache
+    let mut config = SessionConfig::new();
+    config.options_mut().execution.parquet.pushdown_filters = true;
+    config
+        .options_mut()
+        .execution
+        .parquet
+        .max_predicate_cache_size = Some(0);
+    let ctx = SessionContext::new_with_config(config);
+    PredicateCacheTest {
+        // file has 8 rows, which need to be read twice, one for filter, one for
+        // final output
+        expected_inner_records: 16,
+        // Expect this to 0 records read as the cache is disabled. However, it is
+        // non zero due to https://github.com/apache/arrow-rs/issues/8307
+        expected_records: 3,
+    }
+    .run(&ctx)
+    .await
+}
+
+/// Runs the query "SELECT * FROM alltypes_plain WHERE double_col != 0.0"
+/// with a given SessionContext and asserts that the predicate cache metrics
+/// are as expected
+#[derive(Debug)]
+struct PredicateCacheTest {
+    /// Expected records read from the underlying reader (to evaluate filters)
+    /// -- this is the total number of records in the file
+    expected_inner_records: usize,
+    /// Expected records to be read from the cache (after filtering)
+    expected_records: usize,
+}
+
+impl PredicateCacheTest {
+    async fn run(self, ctx: &SessionContext) -> datafusion_common::Result<()> {
+        let Self {
+            expected_inner_records,
+            expected_records,
+        } = self;
+        // Create a dataframe that scans the "alltypes_plain.parquet" file with
+        // a filter on `double_col != 0.0`
+        let path = parquet_test_data() + "/alltypes_plain.parquet";
+        let exec = ctx
+            .read_parquet(path, ParquetReadOptions::default())
+            .await?
+            .filter(col("double_col").not_eq(lit(0.0)))?
+            .create_physical_plan()
+            .await?;
+
+        // run the plan to completion
+        let _ = collect(exec.clone(), ctx.task_ctx()).await?; // run plan
+        let metrics =
+            TestParquetFile::parquet_metrics(&exec).expect("found parquet metrics");
+
+        // verify the predicate cache metrics
+        assert_eq!(
+            get_value(&metrics, "predicate_cache_inner_records"),
+            expected_inner_records
+        );
+        assert_eq!(
+            get_value(&metrics, "predicate_cache_records"),
+            expected_records
+        );
+        Ok(())
+    }
+}

datafusion/datasource-parquet/src/metrics.rs

@@ -72,6 +72,13 @@ pub struct ParquetFileMetrics {
     pub page_index_eval_time: Time,
     /// Total time spent reading and parsing metadata from the footer
     pub metadata_load_time: Time,
+    /// Predicate Cache: number of records read directly from the inner reader.
+    /// This is the number of rows decoded while evaluating predicates
+    pub predicate_cache_inner_records: Count,
+    /// Predicate Cache: number of records read from the cache. This is the
+    /// number of rows that were stored in the cache after evaluating predicates
+    /// reused for the output.
+    pub predicate_cache_records: Count,
 }
 
 impl ParquetFileMetrics {
@@ -140,6 +147,14 @@ impl ParquetFileMetrics {
         let files_ranges_pruned_statistics = MetricBuilder::new(metrics)
             .counter("files_ranges_pruned_statistics", partition);
 
+        let predicate_cache_inner_records = MetricBuilder::new(metrics)
+            .with_new_label("filename", filename.to_string())
+            .counter("predicate_cache_inner_records", partition);
+
+        let predicate_cache_records = MetricBuilder::new(metrics)
+            .with_new_label("filename", filename.to_string())
+            .counter("predicate_cache_records", partition);
+
         Self {
             files_ranges_pruned_statistics,
             predicate_evaluation_errors,
@@ -157,6 +172,8 @@ impl ParquetFileMetrics {
             bloom_filter_eval_time,
             page_index_eval_time,
             metadata_load_time,
+            predicate_cache_inner_records,
+            predicate_cache_records,
         }
     }
 }

datafusion/datasource-parquet/src/opener.rs

@@ -51,6 +51,7 @@ use datafusion_execution::parquet_encryption::EncryptionFactory;
 use futures::{ready, Stream, StreamExt, TryStreamExt};
 use itertools::Itertools;
 use log::debug;
+use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics;
 use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};
 use parquet::arrow::async_reader::AsyncFileReader;
 use parquet::arrow::{ParquetRecordBatchStreamBuilder, ProjectionMask};
@@ -105,6 +106,9 @@ pub(super) struct ParquetOpener {
     #[cfg(feature = "parquet_encryption")]
     pub encryption_factory:
         Option<(Arc<dyn EncryptionFactory>, EncryptionFactoryOptions)>,
+    /// Maximum size of the predicate cache, in bytes. If none, uses
+    /// the arrow-rs default.
+    pub max_predicate_cache_size: Option<usize>,
 }
 
 impl FileOpener for ParquetOpener {
@@ -152,6 +156,7 @@ impl FileOpener for ParquetOpener {
 
         let enable_page_index = self.enable_page_index;
         let encryption_context = self.get_encryption_context();
+        let max_predicate_cache_size = self.max_predicate_cache_size;
 
         Ok(Box::pin(async move {
             let file_decryption_properties = encryption_context
@@ -401,21 +406,42 @@ impl FileOpener for ParquetOpener {
                 builder = builder.with_limit(limit)
             }
 
+            if let Some(max_predicate_cache_size) = max_predicate_cache_size {
+                builder = builder.with_max_predicate_cache_size(max_predicate_cache_size);
+            }
+
+            // metrics from the arrow reader itself
+            let arrow_reader_metrics = ArrowReaderMetrics::enabled();
+
             let stream = builder
                 .with_projection(mask)
                 .with_batch_size(batch_size)
                 .with_row_groups(row_group_indexes)
+                .with_metrics(arrow_reader_metrics.clone())
                 .build()?;
 
-            let stream = stream
-                .map_err(DataFusionError::from)
-                .map(move |b| b.and_then(|b| schema_mapping.map_batch(b)));
+            let files_ranges_pruned_statistics =
+                file_metrics.files_ranges_pruned_statistics.clone();
+            let predicate_cache_inner_records =
+                file_metrics.predicate_cache_inner_records.clone();
+            let predicate_cache_records = file_metrics.predicate_cache_records.clone();
+
+            let stream = stream.map_err(DataFusionError::from).map(move |b| {
+                b.and_then(|b| {
+                    copy_arrow_reader_metrics(
+                        &arrow_reader_metrics,
+                        &predicate_cache_inner_records,
+                        &predicate_cache_records,
+                    );
+                    schema_mapping.map_batch(b)
+                })
+            });
 
             if let Some(file_pruner) = file_pruner {
                 Ok(EarlyStoppingStream::new(
                     stream,
                     file_pruner,
-                    file_metrics.files_ranges_pruned_statistics.clone(),
+                    files_ranges_pruned_statistics,
                 )
                 .boxed())
             } else {
@@ -425,6 +451,22 @@ impl FileOpener for ParquetOpener {
     }
 }
 
+/// Copies metrics from ArrowReaderMetrics (the metrics collected by the
+/// arrow-rs parquet reader) to the parquet file metrics for DataFusion
+fn copy_arrow_reader_metrics(
+    arrow_reader_metrics: &ArrowReaderMetrics,
+    predicate_cache_inner_records: &Count,
+    predicate_cache_records: &Count,
+) {
+    if let Some(v) = arrow_reader_metrics.records_read_from_inner() {
+        predicate_cache_inner_records.add(v);
+    }
+
+    if let Some(v) = arrow_reader_metrics.records_read_from_cache() {
+        predicate_cache_records.add(v);
+    }
+}
+
 /// Wraps an inner RecordBatchStream and a [`FilePruner`]
 ///
 /// This can terminate the scan early when some dynamic filters is updated after
@@ -823,6 +865,7 @@ mod test {
                 expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
+                max_predicate_cache_size: None,
             }
         };
 
@@ -911,6 +954,7 @@ mod test {
                 expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
+                max_predicate_cache_size: None,
             }
         };
 
@@ -1015,6 +1059,7 @@ mod test {
                 expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
+                max_predicate_cache_size: None,
             }
         };
         let make_meta = || FileMeta {
@@ -1129,6 +1174,7 @@ mod test {
                 expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
+                max_predicate_cache_size: None,
             }
         };
 
@@ -1244,6 +1290,7 @@ mod test {
                 expr_adapter_factory: Some(Arc::new(DefaultPhysicalExprAdapterFactory)),
                 #[cfg(feature = "parquet_encryption")]
                 encryption_factory: None,
+                max_predicate_cache_size: None,
             }
         };
 
@@ -1426,6 +1473,7 @@ mod test {
             expr_adapter_factory: None,
             #[cfg(feature = "parquet_encryption")]
             encryption_factory: None,
+            max_predicate_cache_size: None,
         };
 
         let predicate = logical2physical(&col("a").eq(lit(1u64)), &table_schema);

datafusion/datasource-parquet/src/source.rs

@@ -426,6 +426,12 @@ impl ParquetSource {
         self.table_parquet_options.global.bloom_filter_on_read
     }
 
+    /// Return the maxium predicate cache size, in bytes, used when
+    /// `pushdown_filters`
+    pub fn max_predicate_cache_size(&self) -> Option<usize> {
+        self.table_parquet_options.global.max_predicate_cache_size
+    }
+
     /// Applies schema adapter factory from the FileScanConfig if present.
     ///
     /// # Arguments
@@ -580,6 +586,7 @@ impl FileSource for ParquetSource {
             expr_adapter_factory,
             #[cfg(feature = "parquet_encryption")]
             encryption_factory: self.get_encryption_factory_with_config(),
+            max_predicate_cache_size: self.max_predicate_cache_size(),
         })
     }
 

datafusion/proto-common/proto/datafusion_common.proto

@@ -581,6 +581,10 @@ message ParquetOptions {
   oneof coerce_int96_opt {
     string coerce_int96 = 32;
   }
+
+  oneof max_predicate_cache_size_opt {
+    uint64 max_predicate_cache_size = 33;
+  }
 }
 
 enum JoinSide {