cosmetic & fmt

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

Author: aditanase

crates/core/src/delta_datafusion/mod.rs

@@ -58,7 +58,10 @@ use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_expr::logical_plan::CreateExternalTable;
 use datafusion_expr::simplify::SimplifyContext;
 use datafusion_expr::utils::conjunction;
-use datafusion_expr::{col, BinaryExpr, Expr, Extension, LogicalPlan, Operator, TableProviderFilterPushDown, Volatility};
+use datafusion_expr::{
+    col, BinaryExpr, Expr, Extension, LogicalPlan, Operator, TableProviderFilterPushDown,
+    Volatility,
+};
 use datafusion_physical_expr::{create_physical_expr, PhysicalExpr};
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
@@ -867,13 +870,13 @@ pub fn expr_applicable_for_cols(col_names: &[&str], expr: &Expr) -> bool {
         | Expr::IsNotNull(_)
         | Expr::IsNull(_)
         | Expr::Between(_)
-        | Expr::InList(_)  => Ok(TreeNodeRecursion::Continue),
+        | Expr::InList(_) => Ok(TreeNodeRecursion::Continue),
         _ => {
             is_applicable = false;
             Ok(TreeNodeRecursion::Stop)
         }
     })
-        .unwrap();
+    .unwrap();
     is_applicable
 }
 

crates/core/src/kernel/snapshot/log_data.rs

@@ -891,7 +891,7 @@ mod datafusion {
             arrow_cast::cast(batch.column_by_name("output")?, &ArrowDataType::UInt64).ok()
         }
 
-        // This function is required for partition column pruning to be executed correctly
+        // This function is optional but will optimize partition column pruning
         fn contained(&self, column: &Column, value: &HashSet<ScalarValue>) -> Option<BooleanArray> {
             if value.is_empty() || !self.metadata.partition_columns.contains(&column.name) {
                 return None;
@@ -918,9 +918,7 @@ mod datafusion {
                     | ScalarValue::Utf8View(Some(v))
                     | ScalarValue::LargeUtf8(Some(v)) => pv == v,
 
-                    ScalarValue::Dictionary(_, inner) => {
-                        check_scalar(pv, inner)
-                    }
+                    ScalarValue::Dictionary(_, inner) => check_scalar(pv, inner),
                     // FIXME: is this a good enough default or should we sync this with
                     //  expr_applicable_for_cols and bail out with None
                     _ => value.to_string() == pv,
@@ -931,9 +929,11 @@ mod datafusion {
                 if partition_values.is_null(i) {
                     contains.push(false);
                 } else {
-                    contains.push(value.iter().any(|scalar| {
-                        check_scalar(partition_values.value(i), scalar)
-                    }));
+                    contains.push(
+                        value
+                            .iter()
+                            .any(|scalar| check_scalar(partition_values.value(i), scalar)),
+                    );
                 }
             }
 

crates/core/tests/integration_datafusion.rs

@@ -39,17 +39,19 @@ use serial_test::serial;
 use url::Url;
 
 mod local {
+    use super::*;
     use datafusion::datasource::source::DataSourceExec;
     use datafusion::{common::stats::Precision, datasource::provider_as_source};
     use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
-    use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, TableProviderFilterPushDown, TableScan};
+    use datafusion_expr::{
+        LogicalPlan, LogicalPlanBuilder, TableProviderFilterPushDown, TableScan,
+    };
     use deltalake_core::{
         delta_datafusion::DeltaLogicalCodec, logstore::default_logstore, writer::JsonWriter,
     };
     use itertools::Itertools;
     use object_store::local::LocalFileSystem;
     use TableProviderFilterPushDown::{Exact, Inexact};
-    use super::*;
     #[tokio::test]
     #[serial]
     async fn test_datafusion_local() -> TestResult {
@@ -210,26 +212,10 @@ mod local {
         }
     }
 
-    fn find_table_scan(plan: &LogicalPlan) -> Vec<&Expr> {
-        let mut result = vec![];
-
-        plan.apply(|node| {
-            if let LogicalPlan::TableScan(TableScan { ref filters, ..}) = node {
-                result = filters.iter().collect();
-                Ok(TreeNodeRecursion::Stop) // Stop traversal once found
-            } else {
-                Ok(TreeNodeRecursion::Continue) // Continue traversal
-            }
-        }).expect("Traversal should not fail");
-
-        result
-    }
-
     #[tokio::test]
     async fn test_datafusion_query_partitioned_pushdown() -> Result<()> {
         let ctx = SessionContext::new();
-        let table = open_table("../test/tests/data/delta-0.8.0-partitioned")
-            .await?;
+        let table = open_table("../test/tests/data/delta-0.8.0-partitioned").await?;
         ctx.register_table("demo", Arc::new(table.clone()))?;
 
         let pruning_predicates = [
@@ -238,30 +224,40 @@ mod local {
             PruningTestCase::new("year = '2021'"),
             PruningTestCase::with_push_down(
                 "year = '2021' AND day IS NOT NULL",
-                vec![Exact, Exact]
+                vec![Exact, Exact],
             ),
             PruningTestCase::new("year IN ('2021', '2022')"),
             // NOT IN (a, b) is rewritten as (col != a AND col != b)
-            PruningTestCase::with_push_down(
-                "year NOT IN ('2020', '2022')",
-                vec![Exact, Exact]
-            ),
+            PruningTestCase::with_push_down("year NOT IN ('2020', '2022')", vec![Exact, Exact]),
             // BETWEEN a AND b is rewritten as (col >= a AND col < b)
-            PruningTestCase::with_push_down(
-                "year BETWEEN '2021' AND '2022'",
-                vec![Exact, Exact]
-            ),
+            PruningTestCase::with_push_down("year BETWEEN '2021' AND '2022'", vec![Exact, Exact]),
             PruningTestCase::new("year NOT BETWEEN '2019' AND '2020'"),
             PruningTestCase::with_push_down(
                 "year = '2021' AND day IN ('4', '5', '20')",
-                vec![Exact, Exact]
+                vec![Exact, Exact],
             ),
             PruningTestCase::with_push_down(
                 "year = '2021' AND cast(day as int) <= 20",
-                vec![Exact, Inexact]
+                vec![Exact, Inexact],
             ),
         ];
 
+        fn find_scan_filters(plan: &LogicalPlan) -> Vec<&Expr> {
+            let mut result = vec![];
+
+            plan.apply(|node| {
+                if let LogicalPlan::TableScan(TableScan { ref filters, .. }) = node {
+                    result = filters.iter().collect();
+                    Ok(TreeNodeRecursion::Stop) // Stop traversal once found
+                } else {
+                    Ok(TreeNodeRecursion::Continue) // Continue traversal
+                }
+            })
+            .expect("Traversal should not fail");
+
+            result
+        }
+
         for pp in pruning_predicates {
             let pred = pp.sql;
             let sql = format!("SELECT CAST( day as int ) as my_day FROM demo WHERE {pred} ORDER BY CAST( day as int ) ASC");
@@ -271,14 +267,12 @@ mod local {
 
             // validate that we are correctly qualifying filters as Exact or Inexact
             let plan = df.clone().into_optimized_plan()?;
-            let filters = find_table_scan(&plan);
+            let filters = find_scan_filters(&plan);
             let push_down = table.supports_filters_pushdown(&filters)?;
 
             assert_eq!(push_down, pp.push_down);
 
-            let batches = df
-                .collect()
-                .await?;
+            let batches = df.collect().await?;
 
             let batch = &batches[0];