FIXUP - prune DF stats with the same predicate as files

Signed-off-by: Adrian Tanase <[email protected]>

# Conflicts:
#	crates/core/src/delta_datafusion/mod.rs

Author: aditanase

crates/core/src/delta_datafusion/mod.rs

@@ -27,14 +27,17 @@ use std::fmt::{self, Debug};
 use std::sync::Arc;
 
 use arrow_array::types::UInt16Type;
-use arrow_array::{Array, DictionaryArray, RecordBatch, StringArray, TypedDictionaryArray};
+use arrow_array::{
+    Array, BooleanArray, DictionaryArray, RecordBatch, StringArray, TypedDictionaryArray,
+};
 use arrow_cast::display::array_value_to_string;
 use arrow_cast::{cast_with_options, CastOptions};
 use arrow_schema::{
     ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema, SchemaRef,
     SchemaRef as ArrowSchemaRef, TimeUnit,
 };
 use arrow_select::concat::concat_batches;
+use arrow_select::filter::filter_record_batch;
 use async_trait::async_trait;
 use chrono::{DateTime, TimeZone, Utc};
 use datafusion::catalog::{Session, TableProviderFactory};
@@ -87,7 +90,9 @@ use url::Url;
 use crate::delta_datafusion::expr::parse_predicate_expression;
 use crate::delta_datafusion::schema_adapter::DeltaSchemaAdapterFactory;
 use crate::errors::{DeltaResult, DeltaTableError};
-use crate::kernel::{Add, DataCheck, EagerSnapshot, Invariant, Snapshot, StructTypeExt};
+use crate::kernel::{
+    Add, DataCheck, EagerSnapshot, Invariant, LogDataHandler, Snapshot, StructTypeExt,
+};
 use crate::logstore::LogStoreRef;
 use crate::table::builder::ensure_table_uri;
 use crate::table::state::DeltaTableState;
@@ -573,6 +578,8 @@ impl<'a> DeltaScanBuilder<'a> {
                 .unwrap()
         });
 
+        let mut pruning_mask: Option<_> = None;
+
         // Perform Pruning of files to scan
         let (files, files_scanned, files_pruned) = match self.files {
             Some(files) => {
@@ -592,7 +599,9 @@ impl<'a> DeltaScanBuilder<'a> {
                     let files_to_prune = if let Some(predicate) = &logical_filter {
                         let pruning_predicate =
                             PruningPredicate::try_new(predicate.clone(), logical_schema.clone())?;
-                        pruning_predicate.prune(self.snapshot)?
+                        let mask = pruning_predicate.prune(self.snapshot)?;
+                        pruning_mask = Some(mask.clone());
+                        mask
                     } else {
                         vec![true; num_containers]
                     };
@@ -695,10 +704,18 @@ impl<'a> DeltaScanBuilder<'a> {
             ));
         }
 
-        let stats = self
-            .snapshot
-            .datafusion_table_statistics()
-            .unwrap_or(Statistics::new_unknown(&schema));
+        // FIXME - where is the correct place to marry file pruning with statistics pruning?
+        //  Temporarily re-generating the log handler, just so that we can compute the stats.
+        //  Should we update datafusion_table_statistics to optionally take the mask?
+        let stats = if let Some(mask) = pruning_mask {
+            let es = self.snapshot.snapshot();
+            let pruned_stats = prune_file_statistics(&es.files, mask);
+            LogDataHandler::new(&pruned_stats, es.metadata(), es.schema()).statistics()
+        } else {
+            self.snapshot.datafusion_table_statistics()
+        };
+
+        let stats = stats.unwrap_or(Statistics::new_unknown(&schema));
 
         let parquet_options = TableParquetOptions {
             global: self.session.config().options().execution.parquet.clone(),
@@ -756,6 +773,27 @@ impl<'a> DeltaScanBuilder<'a> {
     }
 }
 
+fn prune_file_statistics(
+    record_batches: &Vec<RecordBatch>,
+    pruning_mask: Vec<bool>,
+) -> Vec<RecordBatch> {
+    let mut filtered_batches = Vec::new();
+    let mut mask_offset = 0;
+
+    for batch in record_batches {
+        let num_rows = batch.num_rows();
+        let batch_mask = &pruning_mask[mask_offset..mask_offset + num_rows];
+        mask_offset += num_rows;
+
+        let boolean_mask = BooleanArray::from(batch_mask.to_vec());
+        let filtered_batch =
+            filter_record_batch(batch, &boolean_mask).expect("Failed to filter RecordBatch");
+        filtered_batches.push(filtered_batch);
+    }
+
+    filtered_batches
+}
+
 // TODO: implement this for Snapshot, not for DeltaTable
 #[async_trait]
 impl TableProvider for DeltaTable {

crates/core/tests/integration_datafusion.rs

@@ -41,11 +41,14 @@ use url::Url;
 mod local {
     use super::*;
     use datafusion::datasource::source::DataSourceExec;
+    use datafusion::prelude::SessionConfig;
     use datafusion::{common::stats::Precision, datasource::provider_as_source};
+    use datafusion_common::assert_contains;
     use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
     use datafusion_expr::{
         LogicalPlan, LogicalPlanBuilder, TableProviderFilterPushDown, TableScan,
     };
+    use datafusion_physical_plan::displayable;
     use deltalake_core::{
         delta_datafusion::DeltaLogicalCodec, logstore::default_logstore, writer::JsonWriter,
     };
@@ -212,6 +215,36 @@ mod local {
         }
     }
 
+    #[tokio::test]
+    async fn test_datafusion_optimize_stats_partitioned_pushdown() -> Result<()> {
+        let config = SessionConfig::new().with_target_partitions(2);
+        let ctx = SessionContext::new_with_config(config);
+        let table = open_table("../test/tests/data/http_requests").await?;
+        ctx.register_table("http_requests", Arc::new(table.clone()))?;
+
+        let sql = "SELECT COUNT(*) as num_events FROM http_requests WHERE date > '2023-04-13'";
+        let df = ctx.sql(sql).await?;
+        let plan = df.clone().create_physical_plan().await?;
+
+        // convert to explain plan form
+        let display = displayable(plan.as_ref()).indent(true).to_string();
+
+        assert_contains!(
+            &display,
+            "ProjectionExec: expr=[1437 as num_events]\n  PlaceholderRowExec"
+        );
+
+        let batches = df.collect().await?;
+        let batch = &batches[0];
+
+        assert_eq!(
+            batch.column(0).as_ref(),
+            Arc::new(Int64Array::from(vec![1437])).as_ref(),
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_datafusion_query_partitioned_pushdown() -> Result<()> {
         let ctx = SessionContext::new();
@@ -290,8 +323,7 @@ mod local {
         use datafusion::prelude::*;
         let ctx = SessionContext::new();
         let state = ctx.state();
-        let table = open_table("../test/tests/data/delta-0.8.0")
-            .await?;
+        let table = open_table("../test/tests/data/delta-0.8.0").await?;
 
         // Simple Equality test, we only exercise the limit in this test
         let e = col("value").eq(lit(2));