groupby()
def groupBy(cols: Column*): RelationalGroupedDataset = {
RelationalGroupedDataset(toDF(), cols.map(_.expr), RelationalGroupedDataset.GroupByType)
}
该方法接受一个或多个 Column 对象作为参数,表示要按照哪些列进行分组。在方法内部,首先调用 DataFrame 的 toDF() 方法将当前 DataFrame 转换成一个新的 DataFrame 对象,然后将传入的每个 Column 的表达式加入到一个 Seq[Expression] 中,并将这个序列、转换后的 DataFrame 和一个表示分组类型的枚举值作为参数传递给 RelationalGroupedDataset 的构造函数,最终返回一个 RelationalGroupedDataset 对象。
def groupBy(col1: String, cols: String*): RelationalGroupedDataset = {
val colNames: Seq[String] = col1 +: cols
RelationalGroupedDataset(
toDF(), colNames.map(colName => resolve(colName)), RelationalGroupedDataset.GroupByType)
}
该方法接受一个字符串作为第一个参数,表示要按照哪一列进行分组;接着还接受可变参数,表示其他要按照哪些列进行分组。在方法内部,首先将所有列名存储在一个 Seq[String] 中,然后通过 resolve 方法将列名解析为 Column 对象,最后将转换后的 Column 组成的序列、转换后的 DataFrame 和一个表示分组类型的枚举值作为参数传递给 RelationalGroupedDataset 的构造函数,最终返回一个 RelationalGroupedDataset 对象。
RelationalGroupedDataset 是 Spark SQL 中的一个类,用于对 DataFrame 进行分组操作。RelationalGroupedDataset 提供了一系列操作方法,如 agg, mean, max, min, sum, rollup 等等,这些方法都可以应用于分组数据集上,用于对数据进行聚合操作,并返回结果。
agg()
def agg(exprs: Map[String, String]): DataFrame = groupBy().agg(exprs)
这个函数调用groupBy()函数生成的RelationalGroupedDataset对象,调用其agg()函数
def agg(exprs: Map[String, String]): DataFrame = {
toDF(exprs.map { case (colName, expr) =>
strToExpr(expr)(df(colName).expr)
}.toSeq)
}
这个函数定义在 RelationalGroupedDataset 类中,用于对分组数据集进行聚合操作,返回一个新的 DataFrame 对象。该函数接收一个包含聚合表达式的 Map,其中键表示聚合列名,值表示聚合表达式。函数会应用每个聚合表达式到对应的列上,并返回包含聚合结果的新 DataFrame。
具体解析如下:
- agg(exprs: Map[String, String]): DataFrame 是 RelationalGroupedDataset 类的一个方法。
- exprs 是一个 Map[String, String] 类型参数,表示聚合表达式的集合,其中键为聚合列名,值为聚合表达式的字符串形式。
- 函数体中首先调用 toDF 方法将聚合表达式的结果转化为 DataFrame 对象,并返回。
- toDF 方法的参数是一个序列,包含了每个聚合列的表达式。这里使用了 map 函数将 exprs 中的每个表达式都转化为对应列的表达式。
- 在 map 函数中,首先使用 case 来解构聚合表达式,将列名赋值给 colName,将表达式字符串赋值给 expr。
- 接着,通过 strToExpr(expr)(df(colName).expr) 将表达式字符串转化为 Expression 对象,并应用于对应的列上。这里使用 df(colName).expr 获取列的原始表达式,然后将 expr 转化成 Expression 对象,最后应用到原始表达式上得到新表达式。
最后通过 toSeq 将所有新表达式转化为一个序列,并将该序列作为参数传递给 toDF 方法,生成包含聚合结果的新 DataFrame。
case class Aggregate(
groupingExpressions: Seq[Expression],
aggregateExpressions: Seq[NamedExpression],
child: LogicalPlan)
extends UnaryNode {
override lazy val resolved: Boolean = {
val hasWindowExpressions = aggregateExpressions.exists ( _.collect {
case window: WindowExpression => window
}.nonEmpty
)
expressions.forall(_.resolved) && childrenResolved && !hasWindowExpressions
}
override def output: Seq[Attribute] = aggregateExpressions.map(_.toAttribute)
override def metadataOutput: Seq[Attribute] = Nil
override def maxRows: Option[Long] = {
if (groupingExpressions.isEmpty) {
Some(1L)
} else {
child.maxRows
}
}
final override val nodePatterns : Seq[TreePattern] = Seq(AGGREGATE)
override lazy val validConstraints: ExpressionSet = {
val nonAgg = aggregateExpressions.filter(!_.exists(_.isInstanceOf[AggregateExpression]))
getAllValidConstraints(nonAgg)
}
override protected def withNewChildInternal(newChild: LogicalPlan): Aggregate =
copy(child = newChild)
// Whether this Aggregate operator is group only. For example: SELECT a, a FROM t GROUP BY a
private[sql] def groupOnly: Boolean = {
// aggregateExpressions can be empty through Dateset.agg,
// so we should also check groupingExpressions is non empty
groupingExpressions.nonEmpty && aggregateExpressions.map {
case Alias(child, _) => child
case e => e
}.forall(a => a.foldable || groupingExpressions.exists(g => a.semanticEquals(g)))
}
}
这段代码是 Spark SQL 中 Aggregate 操作的实现代码,是对一个逻辑计划 LogicalPlan 进行聚合操作的节点
private[this] def toDF(aggExprs: Seq[Expression]): DataFrame = {
val aggregates = if (df.sparkSession.sessionState.conf.dataFrameRetainGroupColumns) {
groupingExprs match {
// call `toList` because `Stream` can't serialize in scala 2.13
case s: Stream[Expression] => s.toList ++ aggExprs
case other => other ++ aggExprs
}
} else {
aggExprs
}
val aliasedAgg = aggregates.map(alias)
groupType match {
case RelationalGroupedDataset.GroupByType =>
Dataset.ofRows(df.sparkSession, Aggregate(groupingExprs, aliasedAgg, df.logicalPlan))
case RelationalGroupedDataset.RollupType =>
Dataset.ofRows(
df.sparkSession, Aggregate(Seq(Rollup(groupingExprs.map(Seq(_)))),
aliasedAgg, df.logicalPlan))
case RelationalGroupedDataset.CubeType =>
Dataset.ofRows(
df.sparkSession, Aggregate(Seq(Cube(groupingExprs.map(Seq(_)))),
aliasedAgg, df.logicalPlan))
case RelationalGroupedDataset.PivotType(pivotCol, values) =>
val aliasedGrps = groupingExprs.map(alias)
Dataset.ofRows(
df.sparkSession, Pivot(Some(aliasedGrps), pivotCol, values, aggExprs, df.logicalPlan))
}
}
这段代码定义了一个私有方法 toDF,用于将聚合表达式列表 aggExprs 应用于 DataFrame,生成聚合结果。
Dataset.ofRows(df.sparkSession, Aggregate(groupingExprs, aliasedAgg, df.logicalPlan)) 这行代码将一个聚合操作 Aggregate 应用于 DataFrame,并创建一个新的 Dataset。具体来说,该语句实现了以下功能:
调用 Aggregate 的构造方法,传入三个参数:
- groupingExprs:用于分组的表达式序列。
- aliasedAgg:经过别名处理后的聚合表达式序列。
- df.logicalPlan:当前 DataFrame 的逻辑计划。
再调用 Dataset.ofRows 方法,传入两个参数:
- df.sparkSession:当前 DataFrame 所在的 SparkSession。
- 上一步生成的 Aggregate 对象,作为新 Dataset 的逻辑计划。
由此可以得知,这行代码的作用是生成一个新的 Dataset,其逻辑计划为对原 DataFrame 进行聚合操作的结果