Merge branch 'main' into issue-1925

delta-io · Dec 30, 2023 · 87e5f29 · 87e5f29
2 parents 94b19bf + 74f9d33
commit 87e5f29
Show file tree

Hide file tree

Showing 7 changed files with 907 additions and 82 deletions.
diff --git a/crates/benchmarks/src/bin/merge.rs b/crates/benchmarks/src/bin/merge.rs
@@ -265,6 +265,14 @@ async fn benchmark_merge_tpcds(
         .object_store()
         .delete(&Path::parse("_delta_log/00000000000000000002.json")?)
         .await?;
+    table
+        .object_store()
+        .delete(&Path::parse("_delta_log/00000000000000000003.json")?)
+        .await?;
+    let _ = table
+        .object_store()
+        .delete(&Path::parse("_delta_log/00000000000000000004.json")?)
+        .await;
 
     Ok((duration, metrics))
 }

diff --git a/crates/deltalake-core/src/delta_datafusion/logical.rs b/crates/deltalake-core/src/delta_datafusion/logical.rs
@@ -1,5 +1,7 @@
 //! Logical Operations for DataFusion
 
+use std::collections::HashSet;
+
 use datafusion_expr::{LogicalPlan, UserDefinedLogicalNodeCore};
 
 // Metric Observer is used to update DataFusion metrics from a record batch.
@@ -10,6 +12,7 @@ pub(crate) struct MetricObserver {
     // id is preserved during conversion to physical node
     pub id: String,
     pub input: LogicalPlan,
+    pub enable_pushdown: bool,
 }
 
 impl UserDefinedLogicalNodeCore for MetricObserver {
@@ -35,6 +38,18 @@ impl UserDefinedLogicalNodeCore for MetricObserver {
         write!(f, "MetricObserver id={}", &self.id)
     }
 
+    fn prevent_predicate_push_down_columns(&self) -> HashSet<String> {
+        if self.enable_pushdown {
+            HashSet::new()
+        } else {
+            self.schema()
+                .fields()
+                .iter()
+                .map(|f| f.name().clone())
+                .collect()
+        }
+    }
+
     fn from_template(
         &self,
         _exprs: &[datafusion_expr::Expr],
@@ -43,6 +58,7 @@ impl UserDefinedLogicalNodeCore for MetricObserver {
         MetricObserver {
             id: self.id.clone(),
             input: inputs[0].clone(),
+            enable_pushdown: self.enable_pushdown,
         }
     }
 }
diff --git a/crates/deltalake-core/src/delta_datafusion/mod.rs b/crates/deltalake-core/src/delta_datafusion/mod.rs
@@ -32,7 +32,7 @@ use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema, SchemaR
 use arrow::error::ArrowError;
 use arrow::record_batch::RecordBatch;
 use arrow_array::types::UInt16Type;
-use arrow_array::{Array, DictionaryArray, StringArray};
+use arrow_array::{Array, DictionaryArray, StringArray, TypedDictionaryArray};
 use arrow_cast::display::array_value_to_string;
 
 use arrow_schema::Field;
@@ -132,6 +132,21 @@ fn get_scalar_value(value: Option<&ColumnValueStat>, field: &Arc<Field>) -> Prec
     }
 }
 
+pub(crate) fn get_path_column<'a>(
+    batch: &'a RecordBatch,
+    path_column: &str,
+) -> DeltaResult<TypedDictionaryArray<'a, UInt16Type, StringArray>> {
+    let err = || DeltaTableError::Generic("Unable to obtain Delta-rs path column".to_string());
+    batch
+        .column_by_name(path_column)
+        .unwrap()
+        .as_any()
+        .downcast_ref::<DictionaryArray<UInt16Type>>()
+        .ok_or_else(err)?
+        .downcast_dict::<StringArray>()
+        .ok_or_else(err)
+}
+
 impl DeltaTableState {
     /// Provide table level statistics to Datafusion
     pub fn datafusion_table_statistics(&self) -> DataFusionResult<Statistics> {
@@ -1362,31 +1377,20 @@ fn join_batches_with_add_actions(
 
     let mut files = Vec::with_capacity(batches.iter().map(|batch| batch.num_rows()).sum());
     for batch in batches {
-        let array = batch.column_by_name(path_column).ok_or_else(|| {
-            DeltaTableError::Generic(format!("Unable to find column {}", path_column))
-        })?;
-
-        let iter: Box<dyn Iterator<Item = Option<&str>>> =
-            if dict_array {
-                let array = array
-                    .as_any()
-                    .downcast_ref::<DictionaryArray<UInt16Type>>()
-                    .ok_or(DeltaTableError::Generic(format!(
-                        "Unable to downcast column {}",
-                        path_column
-                    )))?
-                    .downcast_dict::<StringArray>()
-                    .ok_or(DeltaTableError::Generic(format!(
-                        "Unable to downcast column {}",
-                        path_column
-                    )))?;
-                Box::new(array.into_iter())
-            } else {
-                let array = array.as_any().downcast_ref::<StringArray>().ok_or(
-                    DeltaTableError::Generic(format!("Unable to downcast column {}", path_column)),
-                )?;
-                Box::new(array.into_iter())
-            };
+        let err = || DeltaTableError::Generic("Unable to obtain Delta-rs path column".to_string());
+
+        let iter: Box<dyn Iterator<Item = Option<&str>>> = if dict_array {
+            let array = get_path_column(&batch, path_column)?;
+            Box::new(array.into_iter())
+        } else {
+            let array = batch
+                .column_by_name(path_column)
+                .ok_or_else(err)?
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .ok_or_else(err)?;
+            Box::new(array.into_iter())
+        };
 
         for path in iter {
             let path = path.ok_or(DeltaTableError::Generic(format!(