Ray DataFusionDatasource fixes

Author: ccciudatu

src/dataframe.rs

@@ -36,7 +36,7 @@ use datafusion::datasource::memory::DataSourceExec;
 use datafusion::datasource::TableProvider;
 use datafusion::datasource::physical_plan::FileScanConfig;
 use datafusion::datasource::source::DataSource;
-use datafusion::execution::SendableRecordBatchStream;
+use datafusion::execution::{SendableRecordBatchStream};
 use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
 use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 use datafusion::prelude::*;
@@ -708,17 +708,16 @@ impl PyDataFrame {
     }
 
     fn distributed_plan(&self, py: Python<'_>) -> PyResult<DistributedPlan> {
-        let future_plan = self.df.as_ref().clone().create_physical_plan();
-        wait_for_future(py, future_plan)
-            .map(DistributedPlan::new)
-            .map_err(py_datafusion_err)
+        let future_plan = DistributedPlan::try_new(self.df.as_ref());
+        wait_for_future(py, future_plan).map_err(py_datafusion_err)
     }
 
 }
 
 #[pyclass(get_all)]
 #[derive(Debug, Clone)]
 pub struct DistributedPlan {
+    repartition_file_min_size: usize,
     physical_plan: PyExecutionPlan,
 }
 
@@ -755,13 +754,13 @@ impl DistributedPlan {
         }
         let updated_plan = self.plan().clone().transform_up(|node| {
             if let Some(exec) = node.as_any().downcast_ref::<DataSourceExec>() {
-                // Remove redundant ranges from partition files because ParquetExec refuses to repartition
+                // Remove redundant ranges from partition files because FileScanConfig refuses to repartition
                 // if any file has a range defined (even when the range actually covers the entire file).
                 // The EnforceDistribution optimizer rule adds ranges for both full and partial files,
-                // so this tries to rever that to trigger a repartition when no files are actually split.
+                // so this tries to revert that in order to trigger a repartition when no files are actually split.
                 if let Some(file_scan) = exec.data_source().as_any().downcast_ref::<FileScanConfig>() {
-                    let mut file_groups = file_scan.file_groups.clone();
-                    for group in file_groups.iter_mut() {
+                    let mut range_free_file_scan = file_scan.clone();
+                    for group in range_free_file_scan.file_groups.iter_mut() {
                         for file in group.iter_mut() {
                             if let Some(range) = &file.range {
                                 if range.start == 0 && range.end == file.object_meta.size as i64 {
@@ -770,18 +769,14 @@ impl DistributedPlan {
                             }
                         }
                     }
-                    if let Some(repartitioned) = file_scan.clone().with_file_groups(file_groups)
-                        .repartitioned(desired_parallelism, 10 * 1024 * 1024, None)? {
-                        Ok(Transformed::yes(Arc::new(DataSourceExec::new(repartitioned))))
-                    } else {
-                        Ok(Transformed::no(node))
+                    let ordering = range_free_file_scan.eq_properties().output_ordering();
+                    if let Some(repartitioned) = range_free_file_scan
+                        .repartitioned(desired_parallelism, self.repartition_file_min_size, ordering)? {
+                        return Ok(Transformed::yes(Arc::new(DataSourceExec::new(repartitioned))))
                     }
-                } else {
-                    Ok(Transformed::no(node))
                 }
-            } else {
-                Ok(Transformed::no(node))
             }
+            Ok(Transformed::no(node))
         }).map_err(py_datafusion_err)?.data;
         self.physical_plan = PyExecutionPlan::new(updated_plan);
         Ok(())
@@ -790,10 +785,18 @@ impl DistributedPlan {
 
 impl DistributedPlan {
 
-    fn new(plan: Arc<dyn ExecutionPlan>) -> Self {
-        Self {
-            physical_plan: PyExecutionPlan::new(plan)
-        }
+    async fn try_new(df: &DataFrame) -> Result<Self, DataFusionError> {
+        let (mut session_state, logical_plan) = df.clone().into_parts();
+        let repartition_file_min_size = session_state.config_options().optimizer.repartition_file_min_size;
+        // Create the physical plan with a single partition, to ensure that no files are split into ranges.
+        // Otherwise, any subsequent repartition attempt would fail (see the comment in `set_desired_parallelism`)
+        session_state.config_mut().options_mut().execution.target_partitions = 1;
+        let physical_plan = session_state.create_physical_plan(&logical_plan).await?;
+        let physical_plan = PyExecutionPlan::new(physical_plan);
+        Ok(Self {
+            repartition_file_min_size,
+            physical_plan,
+        })
     }
 
     fn plan(&self) -> &Arc<dyn ExecutionPlan> {