python - Pyspark - 循环遍历 structType 和 ArrayType 在 structfield 中进行类型转换

Question

我对 pyspark 很陌生，这个问题让我感到困惑。基本上，我正在寻找一种可扩展方法来通过 structType 或 ArrayType 循环类型转换。

我的数据架构示例:

root
 |-- _id: string (nullable = true)
 |-- created: timestamp (nullable = true)
 |-- card_rates: struct (nullable = true)
 |    |-- rate_1: integer (nullable = true)
 |    |-- rate_2: integer (nullable = true)
 |    |-- rate_3: integer (nullable = true)
 |    |-- card_fee: integer (nullable = true)
 |    |-- payment_method: string (nullable = true)
 |-- online_rates: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- rate_1: integer (nullable = true)
 |    |    |-- rate_2: integer (nullable = true)
 |    |    |-- online_fee: double (nullable = true)
 |-- updated: timestamp (nullable = true)

正如您在这里所看到的，card_rates 是一个结构体，online_rates 是一个结构体数组。我正在寻找循环遍历上面所有字段并有条件地对它们进行类型转换的方法。理想情况下，如果它应该是数字，它应该转换为 double ，如果它应该是字符串，它应该转换为字符串。我需要循环，因为这些 rate_* 字段可能会随着时间的推移而增长。

但现在，我对能够循环它们并将它们全部类型转换为字符串感到满意，因为我对 pyspark 很陌生，并且仍在尝试感受它。

我想要的输出架构:

root
 |-- _id: string (nullable = true)
 |-- created: timestamp (nullable = true)
 |-- card_rates: struct (nullable = true)
 |    |-- rate_1: double (nullable = true)
 |    |-- rate_2: double (nullable = true)
 |    |-- rate_3: double (nullable = true)
 |    |-- card_fee: double (nullable = true)
 |    |-- payment_method: string (nullable = true)
 |-- online_rates: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- rate_1: double (nullable = true)
 |    |    |-- rate_2: double (nullable = true)
 |    |    |-- online_fee: double (nullable = true)
 |-- updated: timestamp (nullable = true)

我已经没有办法做到这一点了。

我从这里得到引用:PySpark convert struct field inside array to string

但是这个解决方案对字段进行了硬编码，并且并没有真正循环字段。

请帮忙。

Answer 1

这是一种借助 StructType.simpleString 和 _parse_datatype_string 内置函数的解决方案:

from pyspark.sql.types import *

df_schema = StructType([
  StructField("_id", StringType(), True),
  StructField("created", TimestampType(), True),
  StructField("card_rates", StructType([
                  StructField("rate_1", IntegerType(), True),
                  StructField("rate_2", IntegerType(), True),
                  StructField("rate_3", IntegerType(), True),
                  StructField("card_fee", IntegerType(), True),
                  StructField("card_fee", IntegerType(), True)])),
  StructField("online_rates", ArrayType(
                  StructType(
                    [
                      StructField("rate_1", IntegerType(),True),
                      StructField("rate_2", IntegerType(),True),
                      StructField("online_fee", DoubleType(),True)
                    ]),True),True),
  StructField("updated", TimestampType(), True)])

schema_str = df_schema.simpleString() # this gives -> struct<_id:string,created:timestamp,card_rates:struct<rate_1:int,rate_2:int,rate_3:int,card_fee:int, card_fee:int>,online_rates:array<struct<rate_1:int,rate_2:int,online_fee:double>>,updated:timestamp>

double_schema = schema_str.replace(':int', ':double')

# convert back to StructType
final_schema = _parse_datatype_string(double_schema)
final_schema

1. 首先使用 schema.simpleString 将架构转换为简单字符串
2. 然后将所有 :int 替换为 :double
3. 最后使用 _parse_datatype_string 将修改后的字符串架构转换为 StructType

更新:

为了避免反引号问题，@jxc 指出更好的解决方案是对元素进行递归扫描，如下所示:

def transform_schema(schema):

  if schema == None:
    return StructType()

  updated = []
  for f in schema.fields:
    if isinstance(f.dataType, IntegerType):
      # if IntegerType convert to DoubleType
      updated.append(StructField(f.name, DoubleType(), f.nullable))
    elif isinstance(f.dataType, ArrayType):
      # if ArrayType unpack the array type(elementType), do recursion then wrap results with ArrayType 
      updated.append(StructField(f.name, ArrayType(transform_schema(f.dataType.elementType))))
    elif isinstance(f.dataType, StructType):
      # if StructType do recursion
      updated.append(StructField(f.name, transform_schema(f.dataType)))
    else:
      # else handle all the other cases i.e TimestampType, StringType etc
      updated.append(StructField(f.name, f.dataType, f.nullable))   

  return StructType(updated)

# call the function with your schema
transform_schema(df_schema)

解释:该函数遍历架构 (StructType) 上的每个项目，并尝试将 int 字段 (StructField) 转换为 double。最后将转换后的schema(StructType)传递给上层(父StructType)。

输出:

StructType(List(
  StructField(_id,StringType,true),
  StructField(created,TimestampType,true),
  StructField(card_rates,
              StructType(List(StructField(rate_1,DoubleType,true),
                              StructField(rate_2,DoubleType,true),
                              StructField(rate_3,DoubleType,true),
                              StructField(card_fee,DoubleType,true),
                              StructField(card_fee,DoubleType,true))),true),
  StructField(online_rates,ArrayType(
    StructType(List(
      StructField(rate_1,DoubleType,true),
      StructField(rate_2,DoubleType,true),
      StructField(online_fee,DoubleType,true))),true),true),
  StructField(updated,TimestampType,true)))

更新:(2020-02-02)

这里是一个关于如何将新架构与现有数据帧一起使用的示例:

updated_schema = transform_schema(df.schema)

# cast each column to the new type
select_expr = [df[f.name].cast(f.dataType) for f in updated_schema.fields]

df.select(*select_expr).show()