Question

我们的设置是PySpark。假设我使用spark.read.csv函数创建数据帧df，即

df = spark.read.csv("directory/name_file.csv")

现在，我需要一种无需手工复制和粘贴即可提取“ name_file”的方法。换句话说，我想要一个仅包含字符串“ name_file”的火花列表或数据框。

请仅提供涉及PySpark SQL或与PySpark兼容的Python代码的解决方案。

问题似乎很简单，但是我花了很多时间寻找解决方案，却一无所获。

Answer 1

有一个功能： input_file_name 然后，您split。

from pyspark.sql import functions as F

df = df.withColumn("path", F.input_file_name())
df = df.withColumn("path_splitted", F.split("path","/"))
df = df.withColumn("name", F.col("path_splitted").getItem(F.size("path_splitted")-1))

df.show()
+---+--------------+--------+----------------+
| id|          path|    name|   path_splitted|
+---+--------------+--------+----------------+
|  1|/foo/bar.csv  |bar.csv |[, foo, bar.csv]|
+---+--------------+--------+----------------+

编辑：在spark 2.4中，您可以使用reverse轻松获取最后一个元素

F.reverse("path_splitted").getItem(0)

Answer 2

如果您不想创建之后需要删除的额外列，您可以链接 pyspark.sql.functions。我们还可以利用 pyspark.sql.functions.element_at (Spark 2.4+)，为我们节省一个操作 (F.size)

df = df.withColumn("filename", F.element_at(F.split(F.input_file_name(), "/"),-1))

或者如果您对父目录名感兴趣

df = df.withColumn("dirname", F.element_at(F.split(F.input_file_name(), "/"),-2))

如何从Spark数据框中提取CSV名称

2 个答案: