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

Question

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

我知道我可以使用schema选项：

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列的数据类型。我怎样才能解决这个问题

Answer 1

根据official documentation，模式可以是StructType或String

我可以给你两个建议：

1-使用虚拟文件的模式

如果您有一个具有相同模式（即一行相同结构）的light文件，则可以将其作为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是一个dict，其（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)

Answer 2

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

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