如何利用PySpark结合Schema DataFrame读取固定宽度文本文件

问题描述 #

现有如下固定宽度格式的文本文件：

00120181120xyz12341
00220180203abc56792
00320181203pqr25483

同时存在一个Schema DataFrame，包含两列：_Name（目标列名）和_Length（对应列的宽度）。需要将上述文本文件转换为如下格式的DataFrame：

#+---+-----+---+
#|C1 |  C2 |C3 |
#+---+-----+---+
#|0  |02018|11 |
#|0  |02018|02 |
#|0  |02018|12 |
#+---+-----+---+

解决方案 #

步骤1：读取固定宽度文本文件 #

先读取文本文件，得到包含原始文本行的DataFrame：

raw_df = spark.read.text("/path/to/your/fixed_width_file.txt")

步骤2：准备Schema DataFrame #

假设你的Schema DataFrame已经存在，这里我们模拟一个符合目标输出的示例（如果是从外部加载，替换为对应的读取逻辑即可）：

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import sum, col, substring

# 初始化SparkSession（如果未初始化）
spark = SparkSession.builder.appName("FixedWidthParser").getOrCreate()

# 模拟Schema DataFrame
schema_df = spark.createDataFrame(
    [("C1", 1), ("C2", 5), ("C3", 2)],
    ["_Name", "_Length"]
)

步骤3：计算字段的起始/结束位置 #

使用窗口函数计算每个字段在固定宽度行中的起始偏移量：

# 定义窗口，按列的顺序（假设Schema DataFrame的顺序就是文本中字段的顺序）
window = Window.orderBy("_Name")

# 计算累计长度，得到每个字段的起始位置（起始位置=前序字段总长度）
schema_with_pos = schema_df.withColumn(
    "start_pos",
    sum("_Length").over(window.rowsBetween(Window.unboundedPreceding, Window.currentRow - 1))
).withColumn(
    "end_pos",
    col("start_pos") + col("_Length")
)

步骤4：动态生成截取表达式 #

遍历Schema的行，生成substring表达式来拆分每个字段：

# 收集Schema的字段信息
schema_fields = schema_with_pos.collect()

# 生成截取表达式
select_expr = []
for field in schema_fields:
    col_name = field["_Name"]
    start = field["start_pos"] + 1  # Spark的substring是1-based索引
    length = field["_Length"]
    select_expr.append(substring(col("value"), start, length).alias(col_name))

步骤5：生成目标DataFrame #

应用表达式并查看结果：

target_df = raw_df.select(*select_expr)
target_df.show()

执行后就会得到你需要的格式的DataFrame。

内容的提问来源于stack exchange，提问作者João Miguens