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

Question

假设我用来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 config 读取数据时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
    }
  }

Answer 1

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

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]带有分区列名称的 。