如何使用Spark确定分区键/列

Posted 2021-05-04

假设我使用partitionBy将一些数据保存到磁盘，例如按日期，所以我的数据看起来像这样：

/mydata/d=01-01-2018/part-00000
/mydata/d=01-01-2018/part-00001
...
/mydata/d=02-01-2018/part-00000
/mydata/d=02-01-2018/part-00001
...

当我使用Hive配置和DataFrame读取数据时，如此

val df = sparkSession.sql(s"select * from $database.$tableName")

我知道：

• 列d上的筛选查询将向下推
• 如果我尝试通过d进行分区（例如GROUP BY d），则不会发生洗牌

但是，假设我不知道分区键是什么（某些上游作业写入数据，并且没有约定）。如何让Spark告诉我哪个是分区键，在这种情况下是d。同样，如果我们有多个分区（例如按月，周，那么一天）。

目前我们拥有的最佳代码非常难看：

def getPartitionColumnsForHiveTable(databaseTableName: String)(implicit sparkSession: SparkSession): Set[String] = {
    val cols = sparkSession.
      sql(s"desc $databaseTableName")
      .select("col_name")
      .collect
      .map(_.getAs[String](0))
      .dropWhile(r => !r.matches("# col_name"))
    if (cols.isEmpty) {
      Set()
    } else {
      cols.tail.toSet
    }
  }

答案

假设您的分区列值中没有=和/，您可以执行以下操作：

val df = spark.sql("show partitions database.test_table")

val partitionedCols: Set[String] = try { 
  df.map(_.getAs[String](0)).first.split('/').map(_.split("=")(0)).toSet
} catch {
  case e: AnalysisException => Set.empty[String]
}

您应该获得带有分区列名称的Array[String]。

另一答案

您可以使用sql语句来获取此信息，show create table <tablename>，describe extended <tablename>或show partitions <tablename>。最后一个给出了最简单的输出解析：

val partitionCols = spark.sql("show partitions <tablename>").as[String].first.split('/').map(_.split("=").head)