[SPARK-30231][SQL][PYTHON] Support explain mode in PySpark df.explain

python/pyspark/sql/dataframe.py

@@ -253,10 +253,18 @@ def printSchema(self):
         print(self._jdf.schema().treeString())
 
     @since(1.3)
-    def explain(self, extended=False):
+    def explain(self, extended=None, mode=None):
         """Prints the (logical and physical) plans to the console for debugging purpose.
 
         :param extended: boolean, default ``False``. If ``False``, prints only the physical plan.
+        :param mode: specifies the expected output format of plans.
+
+            * ``simple``: Print only a physical plan.
+            * ``extended``: Print both logical and physical plans.
+            * ``codegen``: Print a physical plan and generated codes if they are available.
+            * ``cost``: Print a logical plan and statistics if they are available.
+            * ``formatted``: Split explain output into two sections: a physical plan outline \
+                and node details.
 
         >>> df.explain()
         == Physical Plan ==
@@ -271,11 +279,47 @@ def explain(self, extended=False):
         ...
         == Physical Plan ==
         ...
+
+        >>> df.explain(mode="formatted")
+        == Physical Plan ==
+        * Scan ExistingRDD (1)
+        (1) Scan ExistingRDD [codegen id : 1]
+        Output: [age#0, name#1]
+
+        .. versionchanged:: 3.0.0
+           Added optional argument `mode` to specify the expected output format of plans.
         """
-        if extended:
-            print(self._jdf.queryExecution().toString())
-        else:
-            print(self._jdf.queryExecution().simpleString())
+
+        if extended is not None and mode is not None:
+            raise Exception("extended and mode can not be specified simultaneously")
+
+        # For the no argument case: df.explain()
+        is_no_argument = extended is None and mode is None
+
+        # For the cases below:
+        #   explain(True)
+        #   explain(extended=False)
+        is_extended_case = extended is not None and isinstance(extended, bool)
+
+        # For the mode specified: df.explain(mode="formatted")
+        is_mode_case = mode is not None and isinstance(mode, basestring)
+
+        if not is_no_argument and not (is_extended_case or is_mode_case):
+            argtypes = [
+                str(type(arg)) for arg in [extended, mode] if arg is not None]
+            raise TypeError(
+                "extended (optional) and mode (optional) should be a bool and str; "
+                "however, got [%s]." % ", ".join(argtypes))
+
+        # Sets an explain mode depending on a given argument
+        if is_no_argument:
+            explainMode = "simple"
+        elif is_extended_case:
+            explainMode = "extended" if extended else "simple"
+        elif is_mode_case:
+            explainMode = mode
+
+        print(self._jdf.toExplainString(explainMode))
 
     @since(2.4)
     def exceptAll(self, other):

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.sql
 
 import java.io.{ByteArrayOutputStream, CharArrayWriter, DataOutputStream}
+import java.util.Locale
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ArrayBuffer
@@ -521,20 +522,14 @@ class Dataset[T] private[sql](
   def printSchema(level: Int): Unit = println(schema.treeString(level))
   // scalastyle:on println
 
-  /**
-   * Prints the plans (logical and physical) with a format specified by a given explain mode.
-   *
-   * @group basic
-   * @since 3.0.0
-   */
-  def explain(mode: ExplainMode): Unit = {
+  private def toExplainString(mode: ExplainMode): String = {
     // Because temporary views are resolved during analysis when we create a Dataset, and
     // `ExplainCommand` analyzes input query plan and resolves temporary views again. Using
     // `ExplainCommand` here will probably output different query plans, compared to the results
     // of evaluation of the Dataset. So just output QueryExecution's query plans here.
     val qe = ExplainCommandUtil.explainedQueryExecution(sparkSession, logicalPlan, queryExecution)
 
-    val outputString = mode match {
+    mode match {
       case ExplainMode.Simple =>
         qe.simpleString
       case ExplainMode.Extended =>
@@ -550,8 +545,29 @@ class Dataset[T] private[sql](
       case ExplainMode.Formatted =>
         qe.simpleString(formatted = true)
     }
+  }
+
+  private[sql] def toExplainString(mode: String): String = {
+    mode.toLowerCase(Locale.ROOT) match {
+      case "simple" => toExplainString(ExplainMode.Simple)
+      case "extended" => toExplainString(ExplainMode.Extended)
+      case "codegen" => toExplainString(ExplainMode.Codegen)
+      case "cost" => toExplainString(ExplainMode.Cost)
+      case "formatted" => toExplainString(ExplainMode.Formatted)
+      case _ => throw new IllegalArgumentException(s"Unknown explain mode: $mode. Accepted " +
+        "explain modes are 'simple', 'extended', 'codegen', 'cost', 'formatted'.")
+    }
+  }
+
+  /**
+   * Prints the plans (logical and physical) with a format specified by a given explain mode.
+   *
+   * @group basic
+   * @since 3.0.0
+   */
+  def explain(mode: ExplainMode): Unit = {
     // scalastyle:off println
-    println(outputString)
+    println(toExplainString(mode))
     // scalastyle:on println
   }
 

sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala

@@ -309,6 +309,24 @@ class ExplainSuite extends QueryTest with SharedSparkSession {
         "(1) LocalTableScan [codegen id :" ::
         Nil: _*)
   }
+
+  test("Dataset.toExplainString has mode as string") {
+    val df = spark.range(10).toDF
+    def assertExplainOutput(mode: ExplainMode): Unit = {
+      assert(df.toExplainString(mode.toString).replaceAll("#\\d+", "#x").trim ===
+        getNormalizedExplain(df, mode).trim)
+    }
+    assertExplainOutput(ExplainMode.Simple)
+    assertExplainOutput(ExplainMode.Extended)
+    assertExplainOutput(ExplainMode.Codegen)
+    assertExplainOutput(ExplainMode.Cost)
+    assertExplainOutput(ExplainMode.Formatted)
+
+    val errMsg = intercept[IllegalArgumentException] {
+      df.toExplainString("unknown")
+    }.getMessage
+    assert(errMsg.contains("Unknown explain mode: unknown"))
+  }
 }
 
 case class ExplainSingleData(id: Int)