首页 > 解决方案 > rlike 语句时编写巨大的功能方式

问题描述

我正在使用正则表达式根据 DataFrame 中的扩展名来识别文件类型。

import org.apache.spark.sql.{Column, DataFrame}
val ignoreCase        :String = "(?i)"
val ignoreExtension   :String = "(?:\\.[_\\d]+)*(?:|\\.bck|\\.old|\\.orig|\\.bz2|\\.gz|\\.7z|\\.z|\\.zip)*(?:\\.[_\\d]+)*$"
val pictureFileName   :String = "image"
val pictureFileType   :String = ignoreCase + "^.+(?:\\.gif|\\.ico|\\.jpeg|\\.jpg|\\.png|\\.svg|\\.tga|\\.tif|\\.tiff|\\.xmp)" + ignoreExtension
val videoFileName     :String = "video"
val videoFileType     :String = ignoreCase + "^.+(?:\\.mod|\\.mp4|\\.mkv|\\.avi|\\.mpg|\\.mpeg|\\.flv)" + ignoreExtension
val otherFileName     :String = "other"

def pathToExtension(cl: Column): Column = {
    when(cl.rlike( pictureFileType ), pictureFileName ).
    when(cl.rlike( videoFileType ), videoFileName ).
    otherwise(otherFileName)
}

val df = List("file.jpg", "file.avi", "file.jpg", "file3.tIf", "file5.AVI.zip", "file4.mp4","afile" ).toDF("filename")
val df2 = df.withColumn("filetype",  pathToExtension( col( "filename" ) ) )
df2.show

这只是一个示例,我已经识别了 30 个正则表达式和类型,因此该函数pathToExtension()非常长,因为我必须when为每种类型添加一个新语句。
我找不到正确的方法来使用包含正则表达式和名称的列表或映射以功能方式编写此代码,如下所示:

val typelist = List((pictureFileName,pictureFileType),(videoFileName,videoFileType)) 
foreach [need help for this part]  

到目前为止我尝试过的所有代码都无法正常工作。

标签: scalaapache-spark

解决方案


您可以使用foldLeft来遍历您的when条件列表并将它们链接起来,如下所示:

import org.apache.spark.sql.Column
import org.apache.spark.sql.functions._
import spark.implicits._

val default = "other"

def chainedWhen(c: Column, rList: List[(String, String)]): Column = rList.tail.
  foldLeft(when(c rlike rList.head._2, rList.head._1))( (acc, t) =>
    acc.when(c rlike t._2, t._1)
  ).otherwise(default)

测试方法:

val df = Seq(
  (1, "a.txt"), (2, "b.gif"), (3, "c.zip"), (4, "d.oth")
).toDF("id", "file_name")

val rList = List(("text", ".*\\.txt"), ("gif", ".*\\.gif"), ("zip", ".*\\.zip"))

df.withColumn("file_type", chainedWhen($"file_name", rList)).show
// +---+---------+---------+
// | id|file_name|file_type|
// +---+---------+---------+
// |  1|    a.txt|     text|
// |  2|    b.gif|      gif|
// |  3|    c.zip|      zip|
// |  4|    d.oth|    other|
// +---+---------+---------+

推荐阅读