使用枚举以获取来自DataFrame的分区列

Question

我试图将所有列及其数据类型放入一个变量中，也只有将分区列纳入Python列表类型的另一个变量。

从描述扩展中获取细节。

df = spark.sql("describe extended schema_name.table_name")

    +----------------------------------------------------------+
    |col_name                    |data_type                                                                       |
    +----------------------------+-----------------------------+
    |col1                        |string                                                                          |
    |col2                        |int                                                                      
    |col3                        |string   
    |col4                        |int                                                                      
    |col5                        |string   
    |# Partition Information     |                                                                                |
    |# col_name                  |data_type                                                                       |
    |col4                        |int                                                                          |
    |col5                        |string                                                                          |
    |                            |                                                                                |
    |# Detailed Table Information|                                                                                |
    |Database                    |schema_name                                                                         |
    |Table                       |table_name                                                                        |
    |Owner                       |owner.name                                                                  |

将结果转换为列表。

des_list=df.select(df.col_name,df.data_type).rdd.map(lambda x:(x[0],x[1])).collect()

这是我尝试获取所有列（直到＃分区信息之前的所有项目）的方式。

  all_cols_name_type=[]
    for index,item in enumerate(des_list):
        if item[0]=='# Partition Information':
            all_cols_name_type.append(des_list[:index])

    

对于分区，我想在项目“＃col_name”和行之间获得所有内容

。

Answer 1

您可以尝试使用以下答案或在Scala中等效：

    val (partitionCols, dataCols) = spark.catalog.listColumns("schema_name.table_name")
      .collect()
      .partition(c => c.isPartition)

    val parCols = partitionCols.map(c => (c.name, c.dataType))
    val datCols = dataCols.map(c => (c.name, c.dataType))

如果表格未定义（EG Reading Parquet Dataset中）直接使用spark.read.parquet（“ s3：// path/...”）），然后您可以在scala中使用以下片段：

    val (partitionSchema, dataSchema) = df.queryExecution.optimizedPlan match {
      case LogicalRelation(hfs: HadoopFsRelation, _, _, _) =>
       (hfs.partitionSchema, hfs.dataSchema)
      case DataSourceV2ScanRelation(_, scan: FileScan, _) =>
        (scan.readPartitionSchema, scan.readDataSchema)
      case _ => (StructType(Seq()), StructType(Seq()))
    }

    val parCols = partitionSchema.map(f => (f.name, f.dataType))
    val datCols = dataSchema.map(f => (f.name, f.dataType))

Answer 2

这样做有一个技巧：您可以使用 nootonalyy_increasing_id 要给每一行，找到具有＃col_name的行，然后获取该索引。像这样的东西

我的示例表

df = spark.sql('describe data')
df = df.withColumn('id', F.monotonically_increasing_id())
df.show()

+--------------------+---------+-------+---+
|            col_name|data_type|comment| id|
+--------------------+---------+-------+---+
|                  c1|      int|   null|  0|
|                  c2|   string|   null|  1|
|# Partition Infor...|         |       |  2|
|          # col_name|data_type|comment|  3|
|                  c2|   string|   null|  4|
+--------------------+---------+-------+---+

很棘手的部分

idx = df.where(F.col('col_name') == '# col_name').first()['id']
# 3

partition_cols = [r['col_name'] for r in df.where(F.col('id') > idx).collect()]
# ['c2']