[SPARK-40852][CONNECT][PYTHON] Introduce StatFunction in proto and implement DataFrame.summary

connector/connect/src/main/protobuf/spark/connect/relations.proto

@@ -48,6 +48,8 @@ message Relation {
     SubqueryAlias subquery_alias = 16;
     Repartition repartition = 17;
 
+    StatFunction stat_function = 100;
+
     Unknown unknown = 999;
   }
 }
@@ -254,3 +256,21 @@ message Repartition {
   // Optional. Default value is false.
   bool shuffle = 3;
 }
+
+// StatFunction
+message StatFunction {
+  // Required. The input relation.
+  Relation input = 1;
+  // Required. The function and its parameters.
+  oneof function {
+    Summary summary = 2;
+
+    Unknown unknown = 999;
+  }
+
+  // StatFunctions.summary
+  message Summary {
+    repeated string statistics = 1;
+  }
+}
+

connector/connect/src/main/scala/org/apache/spark/sql/connect/dsl/package.scala

@@ -441,6 +441,22 @@ package object dsl {
             Repartition.newBuilder().setInput(logicalPlan).setNumPartitions(num).setShuffle(true))
           .build()
 
+      def summary(statistics: String*): Relation = {
+        Relation
+          .newBuilder()
+          .setStatFunction(
+            proto.StatFunction
+              .newBuilder()
+              .setInput(logicalPlan)
+              .setSummary(
+                proto.StatFunction.Summary
+                  .newBuilder()
+                  .addAllStatistics(statistics.toSeq.asJava)
+                  .build())
+              .build())
+          .build()
+      }
+
       private def createSetOperation(
           left: Relation,
           right: Relation,

connector/connect/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -21,7 +21,7 @@ import scala.annotation.elidable.byName
 import scala.collection.JavaConverters._
 
 import org.apache.spark.connect.proto
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.{Dataset, SparkSession}
 import org.apache.spark.sql.catalyst.AliasIdentifier
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar}
 import org.apache.spark.sql.catalyst.expressions
@@ -32,6 +32,7 @@ import org.apache.spark.sql.catalyst.plans.{logical, FullOuter, Inner, JoinType,
 import org.apache.spark.sql.catalyst.plans.logical.{Deduplicate, Except, Intersect, LogicalPlan, Sample, SubqueryAlias, Union}
 import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
 import org.apache.spark.sql.execution.QueryExecution
+import org.apache.spark.sql.execution.stat.StatFunctions
 import org.apache.spark.sql.types._
 import org.apache.spark.util.Utils
 
@@ -73,6 +74,8 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
       case proto.Relation.RelTypeCase.SUBQUERY_ALIAS =>
         transformSubqueryAlias(rel.getSubqueryAlias)
       case proto.Relation.RelTypeCase.REPARTITION => transformRepartition(rel.getRepartition)
+      case proto.Relation.RelTypeCase.STAT_FUNCTION =>
+        transformStatFunction(rel.getStatFunction)
       case proto.Relation.RelTypeCase.RELTYPE_NOT_SET =>
         throw new IndexOutOfBoundsException("Expected Relation to be set, but is empty.")
       case _ => throw InvalidPlanInput(s"${rel.getUnknown} not supported.")
@@ -124,6 +127,19 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
     logical.Range(start, end, step, numPartitions)
   }
 
+  private def transformStatFunction(rel: proto.StatFunction): LogicalPlan = {
+    val child = transformRelation(rel.getInput)
+
+    rel.getFunctionCase match {
+      case proto.StatFunction.FunctionCase.SUMMARY =>
+        StatFunctions
+          .summary(Dataset.ofRows(session, child), rel.getSummary.getStatisticsList.asScala.toSeq)
+          .logicalPlan
+
+      case _ => throw InvalidPlanInput(s"StatFunction ${rel.getUnknown} not supported.")
+    }
+  }
+
   private def transformDeduplicate(rel: proto.Deduplicate): LogicalPlan = {
     if (!rel.hasInput) {
       throw InvalidPlanInput("Deduplicate needs a plan input")

connector/connect/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectProtoSuite.scala

@@ -261,6 +261,12 @@ class SparkConnectProtoSuite extends PlanTest with SparkConnectPlanTest {
     comparePlans(connectPlan2, sparkPlan2)
   }
 
+  test("Test summary") {
+    comparePlans(
+      connectTestRelation.summary("count", "mean", "stddev"),
+      sparkTestRelation.summary("count", "mean", "stddev"))
+  }
+
   private def createLocalRelationProtoByQualifiedAttributes(
       attrs: Seq[proto.Expression.QualifiedAttribute]): proto.Relation = {
     val localRelationBuilder = proto.LocalRelation.newBuilder()

python/pyspark/sql/connect/dataframe.py

@@ -376,6 +376,16 @@ def unionByName(self, other: "DataFrame", allowMissingColumns: bool = False) ->
     def where(self, condition: Expression) -> "DataFrame":
         return self.filter(condition)
 
+    def summary(self, *statistics: str) -> "DataFrame":
+        _statistics: List[str] = list(statistics)
+        for s in _statistics:
+            if not isinstance(s, str):
+                raise TypeError(f"'statistics' must be list[str], but got {type(s).__name__}")
+        return DataFrame.withPlan(
+            plan.StatFunction(child=self._plan, function="summary", statistics=_statistics),
+            session=self._session,
+        )
+
     def _get_alias(self) -> Optional[str]:
         p = self._plan
         while p is not None:

python/pyspark/sql/connect/plan.py

@@ -16,6 +16,7 @@
 #
 
 from typing import (
+    Any,
     List,
     Optional,
     Sequence,
@@ -750,3 +751,40 @@ def _repr_html_(self) -> str:
             </li>
         </uL>
         """
+
+
+class StatFunction(LogicalPlan):
+    def __init__(self, child: Optional["LogicalPlan"], function: str, **kwargs: Any) -> None:
+        super().__init__(child)
+        assert function in ["summary"]
+        self.function = function
+        self.kwargs = kwargs
+
+    def plan(self, session: Optional["RemoteSparkSession"]) -> proto.Relation:
+        assert self._child is not None
+
+        plan = proto.Relation()
+        plan.stat_function.input.CopyFrom(self._child.plan(session))
+
+        if self.function == "summary":
+            plan.stat_function.summary.statistics.extend(self.kwargs.get("statistics", []))
+        else:
+            raise Exception(f"Unknown function ${self.function}.")
+
+        return plan
+
+    def print(self, indent: int = 0) -> str:
+        i = " " * indent
+        return f"""{i}<StatFunction function='{self.function}' augments='{self.kwargs}'>"""
+
+    def _repr_html_(self) -> str:
+        return f"""
+        <ul>
+           <li>
+              <b>StatFunction</b><br />
+              Function: {self.function} <br />
+              Augments: {self.kwargs} <br />
+              {self._child_repr_()}
+           </li>
+        </ul>
+        """