chore: change variable naming, and update documentation to make clear how the datasource repartitioning is configured and performed

Author: wiedld

datafusion/common/src/config.rs

@@ -632,14 +632,21 @@ config_namespace! {
         /// long runner execution, all types of joins may encounter out-of-memory errors.
         pub allow_symmetric_joins_without_pruning: bool, default = true
 
-        /// When set to `true`, file groups will be repartitioned to achieve maximum parallelism.
-        /// Currently Parquet and CSV formats are supported.
+        /// When set to `true`, datasource partitions will be repartitioned to achieve maximum parallelism.
+        /// This applies to both in-memory partitions and FileSource's file groups (1 group is 1 partition).
         ///
-        /// If set to `true`, all files will be repartitioned evenly (i.e., a single large file
+        /// For FileSources, only Parquet and CSV formats are currently supported.
+        ///
+        /// If set to `true` for a FileSource, all files will be repartitioned evenly (i.e., a single large file
         /// might be partitioned into smaller chunks) for parallel scanning.
-        /// If set to `false`, different files will be read in parallel, but repartitioning won't
+        /// If set to `false` for a FileSource, different files will be read in parallel, but repartitioning won't
         /// happen within a single file.
-        pub repartition_file_scans: bool, default = true
+        ///
+        /// If set to `true` for an in-memory source, all memtable's partitions will have their batches
+        /// repartitioned evenly to the desired number of `target_partitions`. Repartitioning can change
+        /// the total number of partitions and batches per partition, but does not slice the initial
+        /// record tables provided to the MemTable on creation.
+        pub repartition_datasource_scans: bool, default = true
 
         /// Should DataFusion repartition data using the partitions keys to execute window
         /// functions in parallel using the provided `target_partitions` level

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -364,7 +364,7 @@ fn ensure_distribution_helper(
     let mut config = ConfigOptions::new();
     config.execution.target_partitions = target_partitions;
     config.optimizer.enable_round_robin_repartition = true;
-    config.optimizer.repartition_file_scans = false;
+    config.optimizer.repartition_datasource_scans = false;
     config.optimizer.repartition_file_min_size = 1024;
     config.optimizer.prefer_existing_sort = prefer_existing_sort;
     ensure_distribution(distribution_context, &config).map(|item| item.data.plan)
@@ -396,7 +396,7 @@ fn test_suite_default_config_options() -> ConfigOptions {
     config.optimizer.prefer_existing_union = false;
 
     // By default, will not repartition file scans.
-    config.optimizer.repartition_file_scans = false;
+    config.optimizer.repartition_datasource_scans = false;
     config.optimizer.repartition_file_min_size = 1024;
 
     // By default, set query execution concurrency to 10.
@@ -449,7 +449,7 @@ impl TestConfig {
     /// If preferred, will repartition file scans.
     /// Accepts a minimum file size to repartition.
     fn with_prefer_repartition_file_scans(mut self, file_min_size: usize) -> Self {
-        self.config.optimizer.repartition_file_scans = true;
+        self.config.optimizer.repartition_datasource_scans = true;
         self.config.optimizer.repartition_file_min_size = file_min_size;
         self
     }
@@ -3483,7 +3483,7 @@ async fn test_distribute_sort_parquet() -> Result<()> {
     let test_config: TestConfig =
         TestConfig::default().with_prefer_repartition_file_scans(1000);
     assert!(
-        test_config.config.optimizer.repartition_file_scans,
+        test_config.config.optimizer.repartition_datasource_scans,
         "should enable scans to be repartitioned"
     );
 
@@ -3524,7 +3524,7 @@ async fn test_distribute_sort_memtable() -> Result<()> {
     let test_config: TestConfig =
         TestConfig::default().with_prefer_repartition_file_scans(1000);
     assert!(
-        test_config.config.optimizer.repartition_file_scans,
+        test_config.config.optimizer.repartition_datasource_scans,
         "should enable scans to be repartitioned"
     );
 

datafusion/datasource/src/memory.rs

@@ -356,7 +356,9 @@ impl MemoryExec {
 /// Data source configuration for reading in-memory batches of data
 #[derive(Clone, Debug)]
 pub struct MemorySourceConfig {
-    /// The partitions to query
+    /// The partitions to query.
+    ///
+    /// Each partition is a `Vec<RecordBatch>`.
     partitions: Vec<Vec<RecordBatch>>,
     /// Schema representing the data before projection
     schema: SchemaRef,
@@ -448,14 +450,18 @@ impl DataSource for MemorySourceConfig {
         }
     }
 
+    /// If possible, redistribute batches across partitions according to their size.
+    ///
+    /// Returns `Ok(None)` if unable to repartition. Preserve output ordering if exists.
+    /// Refer to [`DataSource::repartitioned`] for further details.
     fn repartitioned(
         &self,
         target_partitions: usize,
         _repartition_file_min_size: usize,
         output_ordering: Option<LexOrdering>,
     ) -> Result<Option<Arc<dyn DataSource>>> {
         if self.partitions.is_empty() || self.partitions.len() >= target_partitions
-        // if already have more partitions than desired, do not merge
+        // if have no partitions, or already have more partitions than desired, do not repartition
         {
             return Ok(None);
         }
@@ -758,7 +764,9 @@ impl MemorySourceConfig {
 
     /// Repartition while preserving order.
     ///
-    /// Returns None if cannot fulfill requested repartitioning.
+    /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+    /// as having too few batches to fulfill the `target_partitions` or if unable
+    /// to preserve output ordering.
     fn repartition_preserving_order(
         &self,
         target_partitions: usize,
@@ -778,6 +786,8 @@ impl MemorySourceConfig {
 
             let cnt_to_repartition = target_partitions - self.partitions.len();
 
+            // Label the current partitions and their order.
+            // Such that when we later split up the partitions into smaller sizes, we are maintaining the order.
             let to_repartition = self
                 .partitions
                 .iter()
@@ -789,11 +799,15 @@ impl MemorySourceConfig {
                 })
                 .collect_vec();
 
-            // split the largest partitions
+            // Put all of the partitions into a heap ordered by `RePartition::partial_cmp`, which sizes
+            // by count of rows.
             let mut max_heap = BinaryHeap::with_capacity(target_partitions);
             for rep in to_repartition {
                 max_heap.push(rep);
             }
+
+            // Split the largest partitions into smaller partitions. Maintaining the output
+            // order of the partitions & newly created partitions.
             let mut cannot_split_further = Vec::with_capacity(target_partitions);
             for _ in 0..cnt_to_repartition {
                 loop {
@@ -814,6 +828,8 @@ impl MemorySourceConfig {
             }
             let mut partitions = max_heap.drain().collect_vec();
             partitions.extend(cannot_split_further);
+
+            // Finally, sort all partitions by the output ordering.
             partitions.sort_by_key(|p| p.idx);
             let partitions = partitions.into_iter().map(|rep| rep.batches).collect_vec();
 
@@ -824,7 +840,8 @@ impl MemorySourceConfig {
     /// Repartition into evenly sized chunks (as much as possible without batch splitting),
     /// disregarding any ordering.
     ///
-    /// Returns None if cannot fulfill requested repartitioning.
+    /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+    /// as having too few batches to fulfill the `target_partitions`.
     fn repartition_evenly_by_size(
         &self,
         target_partitions: usize,
@@ -835,7 +852,7 @@ impl MemorySourceConfig {
             return Ok(None);
         }
 
-        // repartition evenly
+        // Take all flattened batches (all in 1 partititon/vec) and divide evenly into the desired number of `target_partitions`.
         let total_num_rows = flatten_batches.iter().map(|b| b.num_rows()).sum::<usize>();
         let target_partition_size = total_num_rows.div_ceil(target_partitions);
         let mut partitions =
@@ -871,8 +888,12 @@ impl MemorySourceConfig {
 ///
 /// Do not implement clone, in order to avoid unnecessary copying during repartitioning.
 struct RePartition {
+    /// Original output ordering for the partition.
     idx: usize,
+    /// Total size of the partition, for use in heap ordering
+    /// (a.k.a. splitting up the largest partitions).
     row_count: usize,
+    /// A partition containing record batches.
     batches: Vec<RecordBatch>,
 }
 
@@ -886,12 +907,12 @@ impl RePartition {
         }
 
         let new_0 = RePartition {
-            idx: self.idx,
+            idx: self.idx, // output ordering
             row_count: 0,
             batches: vec![],
         };
         let new_1 = RePartition {
-            idx: self.idx + 1,
+            idx: self.idx + 1, // output ordering +1
             row_count: 0,
             batches: vec![],
         };

datafusion/datasource/src/source.rs

@@ -60,7 +60,15 @@ pub trait DataSource: Send + Sync + Debug {
     /// Format this source for display in explain plans
     fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result;
 
-    /// Return a copy of this DataSource with a new partitioning scheme
+    /// Return a copy of this DataSource with a new partitioning scheme.
+    ///
+    /// Returns `Ok(None)` (the default) if the partitioning cannot be changed.
+    /// Refer to [`ExecutionPlan::repartitioned`] for details on when None should be returned.
+    ///
+    /// Repartitioning should not change the output ordering, if this ordering exists.
+    /// Refer to [`MemorySourceConfig::repartitioned`](crate::memory::MemorySourceConfig) and the FileSource's
+    /// [`FileGroupPartitioner::repartition_file_groups`](crate::file_groups::FileGroupPartitioner::repartition_file_groups)
+    /// for examples.
     fn repartitioned(
         &self,
         _target_partitions: usize,
@@ -148,6 +156,10 @@ impl ExecutionPlan for DataSourceExec {
         Ok(self)
     }
 
+    /// Implementation of [`ExecutionPlan::repartitioned`] which relies upon the inner [`DataSource::repartitioned`].
+    ///
+    /// If the data source does not support changing its partitioning, returns `Ok(None)` (the default). Refer
+    /// to [`ExecutionPlan::repartitioned`] for more details.
     fn repartitioned(
         &self,
         target_partitions: usize,

datafusion/execution/src/config.rs

@@ -307,7 +307,7 @@ impl SessionConfig {
 
     /// Enables or disables the use of repartitioning for file scans
     pub fn with_repartition_file_scans(mut self, enabled: bool) -> Self {
-        self.options.optimizer.repartition_file_scans = enabled;
+        self.options.optimizer.repartition_datasource_scans = enabled;
         self
     }
 

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -1155,6 +1155,10 @@ fn get_repartition_requirement_status(
 /// operators to satisfy distribution requirements. Since this function
 /// takes care of such requirements, we should avoid manually adding data
 /// exchange operators in other places.
+///
+/// This function is intended to be used in a bottom up traversal, as it
+/// can first repartition (or newly partition) at the datasources -- these
+/// source partitions may be later repartitioned with additional data exchange operators.
 pub fn ensure_distribution(
     dist_context: DistributionContext,
     config: &ConfigOptions,
@@ -1168,7 +1172,7 @@ pub fn ensure_distribution(
     let target_partitions = config.execution.target_partitions;
     // When `false`, round robin repartition will not be added to increase parallelism
     let enable_round_robin = config.optimizer.enable_round_robin_repartition;
-    let repartition_file_scans = config.optimizer.repartition_file_scans;
+    let repartition_datasource_scans = config.optimizer.repartition_datasource_scans;
     let batch_size = config.execution.batch_size;
     let should_use_estimates = config
         .execution
@@ -1242,9 +1246,13 @@ pub fn ensure_distribution(
                 // Unless partitioning increases the partition count, it is not beneficial:
                 && child.plan.output_partitioning().partition_count() < target_partitions;
 
-            // When `repartition_file_scans` is set, attempt to increase
+            // When `repartition_datasource_scans` is set, attempt to increase
             // parallelism at the source.
-            if repartition_file_scans && roundrobin_beneficial_stats {
+            //
+            // If repartitioning is not possible (a.k.a. None is returned from `ExecutionPlan::repartitioned`)
+            // then no repartitioning will have occurred. As the default implementation returns None, it is only
+            // specific physical plan nodes, such as certain datasources, which are repartitioned.
+            if repartition_datasource_scans && roundrobin_beneficial_stats {
                 if let Some(new_child) =
                     child.plan.repartitioned(target_partitions, config)?
                 {