(spark) branch master updated: [SPARK-52535][SQL] Improve code readability of rule ApplyColumnarRulesAndInsertTransitions

[we...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new ace03f419868 [SPARK-52535][SQL] Improve code readability of rule ApplyColumnarRulesAndInsertTransitions
ace03f419868 is described below

commit ace03f4198687adeb5fbd07a0cb9a6cdb75a46dc
Author: Hongze Zhang <ho...@gmail.com>
AuthorDate: Thu Jul 3 16:30:01 2025 +0800

    [SPARK-52535][SQL] Improve code readability of rule ApplyColumnarRulesAndInsertTransitions
    
    Issue: https://issues.apache.org/jira/browse/SPARK-52535
    
    ### What changes were proposed in this pull request?
    
    A minor cleanup on the code of `ApplyColumnarRulesAndInsertTransitions` to improve readability.
    
    ### Why are the changes needed?
    
    The rule is important for 3rd SQL plugins and the code got messed up a bit during the development iterations back then.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Existing tests.
    
    Closes #51227 from zhztheplayer/wip-transition-rule-cleanup.
    
    Authored-by: Hongze Zhang <ho...@gmail.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../org/apache/spark/sql/execution/Columnar.scala  | 37 ++++++++++++++--------
 1 file changed, 24 insertions(+), 13 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala
index 85ed1a24fc6e..f9193cd0495f 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala
@@ -496,33 +496,31 @@ case class ApplyColumnarRulesAndInsertTransitions(
   extends Rule[SparkPlan] {
 
   /**
-   * Inserts an transition to columnar formatted data.
+   * Ensures columnar output on the input query plan. Transitions will be inserted
+   * on demand.
    */
-  private def insertRowToColumnar(plan: SparkPlan): SparkPlan = {
+  private def ensureOutputsColumnar(plan: SparkPlan): SparkPlan = {
     if (!plan.supportsColumnar) {
       // The tree feels kind of backwards
       // Columnar Processing will start here, so transition from row to columnar
-      RowToColumnarExec(insertTransitions(plan, outputsColumnar = false))
+      RowToColumnarExec(ensureOutputsRowBased(plan))
     } else if (!plan.isInstanceOf[RowToColumnarTransition]) {
-      plan.withNewChildren(plan.children.map(insertRowToColumnar))
+      plan.withNewChildren(plan.children.map(ensureOutputsColumnar))
     } else {
       plan
     }
   }
 
   /**
-   * Inserts RowToColumnarExecs and ColumnarToRowExecs where needed.
+   * Ensures row-based output on the input query plan. Transitions will be inserted
+   * on demand.
    */
-  private def insertTransitions(plan: SparkPlan, outputsColumnar: Boolean): SparkPlan = {
-    if (outputsColumnar) {
-      insertRowToColumnar(plan)
-    } else if (plan.supportsColumnar && !plan.supportsRowBased) {
+  private def ensureOutputsRowBased(plan: SparkPlan): SparkPlan = {
+    if (plan.supportsColumnar && !plan.supportsRowBased) {
       // `outputsColumnar` is false but the plan only outputs columnar format, so add a
       // to-row transition here.
-      ColumnarToRowExec(insertRowToColumnar(plan))
-    } else if (plan.isInstanceOf[ColumnarToRowTransition]) {
-      plan
-    } else {
+      ColumnarToRowExec(ensureOutputsColumnar(plan))
+    } else if (!plan.isInstanceOf[ColumnarToRowTransition]) {
       val outputsColumnar = plan match {
         // With planned write, the write command invokes child plan's `executeWrite` which is
         // neither columnar nor row-based.
@@ -537,6 +535,19 @@ case class ApplyColumnarRulesAndInsertTransitions(
           false
       }
       plan.withNewChildren(plan.children.map(insertTransitions(_, outputsColumnar)))
+    } else {
+      plan
+    }
+  }
+
+  /**
+   * Inserts RowToColumnarExecs and ColumnarToRowExecs where needed.
+   */
+  private def insertTransitions(plan: SparkPlan, outputsColumnar: Boolean): SparkPlan = {
+    if (outputsColumnar) {
+      ensureOutputsColumnar(plan)
+    } else {
+      ensureOutputsRowBased(plan)
     }
   }
 


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org