在spark中使用udf完成包含判断 VS column.isin(XXX:_*)

在这之前，当遇到需要同时判断多个包含关系的时候，我都比较倾向于使用第一种。理由很简单，一个udf，解决所有问题。

但是，今天遇到了一个性能问题（平时肯定也出现了，但是没怎么关注），才关注到这样的现象。

 val filterUsers = (episodesBC: Broadcast[Set[Long]], albumsBC: Broadcast[Set[Long]], channelsBC: Broadcast[Set[Long]]) => udf[Boolean, Long, Long, Long] {
    case (episodeId, albumId, channelId) => {
      if (episodesBC.value.contains(episodeId) || albumsBC.value.contains(albumId) || channelsBC.value.contains(channelId)) true else false
    }
  }


val sql =
      s"""
         |SELECT
         |  episode_id, album_id, channel_id, user_id, cnt
         |FROM
         |  xxx_table
         |WHERE
         |  dt >= "2019-08-25" AND
         |  dt <= "2019-08-30" AND
         |  platform_id in ('ios')
      """.stripMargin


val data = spark.sql(sql)
.withColumn("is_needed", filterUsers((episodeIdBC, albumIdBC, channelBC)(col("episode_id"), col("album_id"), col("channel_id")))
.filter(col("is_needed"))
.select("user_id")
.distinct

data.persist()
val size = data.count()

val splitData = data.randomSplit(Array(0.3, 0.7))

 

 

然而上述代码在读取hive表之后，在data.count的时候，并没有persist。导致数据多次从hive中读取数据，在数据量大的情况下，这非常的糟糕。

查阅了一些资料，没有得到相关的解答。后续若有官方解答，再进行同步。

不能只发现问题，必须KO它，因此开始逐步定位，最终发现，这个udf ：filterUsers 的使用，直接影响了DAG图是否执行persist。

将val data = ....部分代码进行了调整之后，persist生效了。

val data = spark.sql(sql)
.withColumn("episode_is_needed", col("episode_id").isin(episodeIdBC.value:_*))
.withColumn("album_is_needed", col("album_id").isin(albumIdBC.value:_*))
.withColumn("channel_is_needed", col("channel_id").isin(channelBC.value:_*))
.filter(col("episode_is_needed") || col("album_is_needed") || col("channel_is_needed"))
.select("user_id")
.distinct