使用空元素在 pyspark 数据帧 read.csv 中设置架构

Question

我有一个数据集（示例），当导入时

df = spark.read.csv(filename, header=True, inferSchema=True)
df.show()

会将带有“NA”的列分配为 stringType()，我希望它是 IntegerType()（或 ByteType()）。

然后我尝试设置

schema = StructType([
    StructField("col_01", IntegerType()),
    StructField("col_02", DateType()),
    StructField("col_03", IntegerType())
])
df = spark.read.csv(filename, header=True, schema=schema)
df.show()

输出显示'col_03' = null的整行为空。

但是col_01和col_02返回适当的数据，如果它们被调用

df.select(['col_01','col_02']).show()

我可以通过后期转换col_3的数据类型来找到解决此问题的方法

df = spark.read.csv(filename, header=True, inferSchema=True)
df = df.withColumn('col_3',df['col_3'].cast(IntegerType()))
df.show()

，但我认为这并不理想，如果我可以通过设置模式直接为每列分配数据类型会更好。

有人能指导我做错什么吗？或者在导入后转换数据类型是唯一的解决方案？也欢迎对这两种方法的性能提出任何意见（如果我们可以使分配模式起作用）。

谢谢，

Answer 1

您可以使用nullValue以下命令在 spark 的 csv 加载程序中设置新的空值：

对于如下所示的 csv 文件：

col_01,col_02,col_03
111,2007-11-18,3
112,2002-12-03,4
113,2007-02-14,5
114,2003-04-16,NA
115,2011-08-24,2
116,2003-05-03,3
117,2001-06-11,4
118,2004-05-06,NA
119,2012-03-25,5
120,2006-10-13,4

并强制模式：

from pyspark.sql.types import StructType, IntegerType, DateType

schema = StructType([
    StructField("col_01", IntegerType()),
    StructField("col_02", DateType()),
    StructField("col_03", IntegerType())
])

你会得到：

df = spark.read.csv(filename, header=True, nullValue='NA', schema=schema)
df.show()
df.printSchema()

    +------+----------+------+
    |col_01|    col_02|col_03|
    +------+----------+------+
    |   111|2007-11-18|     3|
    |   112|2002-12-03|     4|
    |   113|2007-02-14|     5|
    |   114|2003-04-16|  null|
    |   115|2011-08-24|     2|
    |   116|2003-05-03|     3|
    |   117|2001-06-11|     4|
    |   118|2004-05-06|  null|
    |   119|2012-03-25|     5|
    |   120|2006-10-13|     4|
    +------+----------+------+

    root
     |-- col_01: integer (nullable = true)
     |-- col_02: date (nullable = true)
     |-- col_03: integer (nullable = true)