读取json时预定义数据帧的数据类型

Posted 2023-04-17

【中文标题】读取json时预定义数据帧的数据类型

【英文标题】：Pre-define datatype of dataframe while reading json

【发布时间】：2019-06-17 12:06:17

【问题描述】：

我有一个包含 100 列的 json 文件，我想读取所有列以及两列的预定义数据类型。

我知道我可以使用架构选项来做到这一点：

struct1 = StructType([StructField("npi", StringType(), True), StructField("NCPDP", StringType(), True)

spark.read.json(path=abc.json, schema=struct1)

但是，这段代码只读取了两列：

>>> df.printSchema()
root
 |-- npi: string (nullable = true)
 |-- NCPDP: string (nullable = true)

要使用上面的代码，我必须给出所有 100 列的数据类型。我该如何解决这个问题？

【参考方案1】：

根据official documentation，架构可以是StructType 或String。

我可以建议你 2 个解决方案：

1 - 您使用虚拟文件的架构

如果您有一个具有相同架构的轻文件（即一行相同的结构），您可以将其读取为 Dataframe，然后将架构用于您的其他 json 文件：

df = spark.read.json("/path/to/dummy/file.json")
schm = df.schema
df = spark.read.json(path="abc.json", schema=schm)

2 - 您生成架构

此步骤需要您提供列名（可能还需要提供类型）。 让我们假设col 是一个字典，其中 (key, value) 为 (column name, column type)。

col_list = ['col_name col_type'.format(
    col_name=col_name,
    col_type=col_type,
) for col_name, col_type in col.items()]
schema_string = ', '.join(col_list)
df = spark.read.json(path="abc.json", schema=schema_string)

【参考方案2】：

您可以先读取所有数据，然后转换有问题的两列：

df = spark.read.json(path=abc.json)
df.withColumn("npi", df["npi"].cast("string"))\
  .withColumn("NCPDP", df["NCPDP"].cast("string"))

【讨论】：

对不起，但我错过了提到我有一些前面有零的列，所以我在阅读 df 后无法进行类型转换。