From 17698432f959c483e7158d8bcb3a6a511d5c9675 Mon Sep 17 00:00:00 2001
From: yangzhong <yangzhong@ebay.com>
Date: Fri, 26 May 2023 16:11:04 +0800
Subject: [PATCH] Continue PR 4757

---
 datafusion/execution/src/config.rs            | 24 +++++++++++++++++++
 docs/source/user-guide/cli.md                 |  2 --
 .../user-guide/sql/information_schema.md      |  1 -
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/datafusion/execution/src/config.rs b/datafusion/execution/src/config.rs
index 2867e7bd7ce8d..97770eb99c581 100644
--- a/datafusion/execution/src/config.rs
+++ b/datafusion/execution/src/config.rs
@@ -250,6 +250,30 @@ impl SessionConfig {
         self
     }
 
+    /// Enables or disables the coalescence of small batches into larger batches
+    pub fn with_coalesce_batches(mut self, enabled: bool) -> Self {
+        self.options.execution.coalesce_batches = enabled;
+        self
+    }
+
+    /// Returns true if record batches will be examined between each operator
+    /// and small batches will be coalesced into larger batches.
+    pub fn coalesce_batches(&self) -> bool {
+        self.options.execution.coalesce_batches
+    }
+
+    /// Enables or disables the round robin repartition for increasing parallelism
+    pub fn with_round_robin_repartition(mut self, enabled: bool) -> Self {
+        self.options.optimizer.enable_round_robin_repartition = enabled;
+        self
+    }
+
+    /// Returns true if the physical plan optimizer will try to
+    /// add round robin repartition to increase parallelism to leverage more CPU cores.
+    pub fn round_robin_repartition(&self) -> bool {
+        self.options.optimizer.enable_round_robin_repartition
+    }
+
     /// Convert configuration options to name-value pairs with values
     /// converted to strings.
     ///
diff --git a/docs/source/user-guide/cli.md b/docs/source/user-guide/cli.md
index ec0a917c27f5d..53cceb8d0af1c 100644
--- a/docs/source/user-guide/cli.md
+++ b/docs/source/user-guide/cli.md
@@ -358,7 +358,6 @@ Available commands inside DataFusion CLI are:
 +-------------------------------------------------+---------+
 | datafusion.execution.batch_size                 | 8192    |
 | datafusion.execution.coalesce_batches           | true    |
-| datafusion.execution.coalesce_target_batch_size | 4096    |
 | datafusion.execution.time_zone                  | UTC     |
 | datafusion.explain.logical_plan_only            | false   |
 | datafusion.explain.physical_plan_only           | false   |
@@ -395,7 +394,6 @@ DataFusion CLI v12.0.0
 +-------------------------------------------------+---------+
 | datafusion.execution.batch_size                 | 1024    |
 | datafusion.execution.coalesce_batches           | true    |
-| datafusion.execution.coalesce_target_batch_size | 4096    |
 | datafusion.execution.time_zone                  | UTC     |
 | datafusion.explain.logical_plan_only            | false   |
 | datafusion.explain.physical_plan_only           | false   |
diff --git a/docs/source/user-guide/sql/information_schema.md b/docs/source/user-guide/sql/information_schema.md
index b3fcc843bd9f5..ced6117349d41 100644
--- a/docs/source/user-guide/sql/information_schema.md
+++ b/docs/source/user-guide/sql/information_schema.md
@@ -62,7 +62,6 @@ To show the current session configuration options, use the `SHOW ALL` command or
 +-------------------------------------------------+---------+
 | datafusion.execution.batch_size                 | 8192    |
 | datafusion.execution.coalesce_batches           | true    |
-| datafusion.execution.coalesce_target_batch_size | 4096    |
 | datafusion.execution.time_zone                  | UTC     |
 | datafusion.explain.logical_plan_only            | false   |
 | datafusion.explain.physical_plan_only           | false   |