Implement repartitioned for DeltaScan

delta-io · Apr 15, 2024 · 906e51d · 906e51d
1 parent aa8f4d5
commit 906e51d
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/crates/core/src/delta_datafusion/mod.rs b/crates/core/src/delta_datafusion/mod.rs
@@ -61,6 +61,7 @@ use datafusion_common::scalar::ScalarValue;
 use datafusion_common::tree_node::{TreeNode, TreeNodeVisitor, VisitRecursion};
 use datafusion_common::{
     Column, DFSchema, DataFusionError, Result as DataFusionResult, ToDFSchema,
+    config::ConfigOptions,
 };
 use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::logical_plan::CreateExternalTable;
@@ -860,6 +861,19 @@ impl ExecutionPlan for DeltaScan {
     fn statistics(&self) -> DataFusionResult<Statistics> {
         self.parquet_scan.statistics()
     }
+
+    fn repartitioned(&self, target_partitions: usize, config: &ConfigOptions) -> DataFusionResult<Option<Arc<dyn ExecutionPlan>>> {
+        if let Some(parquet_scan) = self.parquet_scan.repartitioned(target_partitions, config)? {
+            Ok(Some(Arc::new(DeltaScan {
+                table_uri: self.table_uri.clone(),
+                config: self.config.clone(),
+                parquet_scan,
+                logical_schema: self.logical_schema.clone(),
+            })))
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 pub(crate) fn get_null_of_arrow_type(t: &ArrowDataType) -> DeltaResult<ScalarValue> {