从DataFrame的字符串列中删除后缀

Question

我正在使用Spark 2.2，并且想知道在性能方面，从数据框的列中删除后缀的最佳选项是什么。

可以使用UDF函数。

val removeSuffix = udf { (id: String) =>
    if (id != null && id.endsWith("XXX")) {
      id.dropRight(3)
    } else {
      id
    }
  }
df.withColumn("c", udf("col"))

可以使用内置的regexp函数。

df.withColumn("c", regexp_replace($"col", "XXX$", ""))

我知道 udf 通常被认为很慢，但是每一行评估一个正则表达式会更快吗？ [2018-01-21 基于 user8983815 的回答进行更新] 我已经编写了一个基准测试，并且结果有点令人惊讶。

[info] Benchmark                                      Mode  Cnt    Score   Error  Units
[info] RemoveSuffixBenchmark.builtin_optimized        avgt   10  103,188 ± 3,526  ms/op
[info] RemoveSuffixBenchmark.builtin_regexp_replace_  avgt   10   99,173 ± 7,313  ms/op
[info] RemoveSuffixBenchmark.udf                      avgt   10   94,570 ± 5,707  ms/op

对于那些感兴趣的人，代码在这里：https://github.com/YannMoisan/spark-jmh

- Yann Moisan

Answer 1

我认为使用regexp_replace不会导致严重的性能问题，但如果你真的很担心，可以考虑以下几点：

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

def removeSuffix(c: Column) = when(c.endsWith("XXX"), c.substr(lit(0), length(c) - 3)).otherwise(c)

用途：

scala> Seq("fooXXX", "bar").toDF("s").select(removeSuffix($"s").alias("s")).show
+---+
|  s|
+---+
|foo|
|bar|
+---+

- user8983815

Answer 2

如果你将其转换为 DataSet 并执行以下操作呢：

例如：

case class InputData(col1: String, col2: String)
case class OutputData(col1: String, col2: String)

val suffix = "."

def clean(i: InputData): OutputData = {
  OutputData(
    col1 = i.col1.stripSuffix(suffix).toLowerCase,
    col2 = i.col2,
  )
}

df.as[InputData].map(clean).toDF

- Dariusz Krynicki