使用 Spark 引擎执行 SQL 时，如何在 hive UDF 中获取 Spark 的 partitionId 或 taskContext？答案

【问题标题】：How to get the partitionId or taskContext of Spark in hive UDF when SQL executed with Spark engine?使用 Spark 引擎执行 SQL 时，如何在 hive UDF 中获取 Spark 的 partitionId 或 taskContext？
【发布时间】：2021-09-05 18:06:26
【问题描述】：

比如我们用Spark引擎执行下面的SQL，我们需要my_udf(row) 返回 Spark 中的分区 id。

add jar hdfs:///dir/udf/udf.jar; 
create temporary function my_udf as 'com.my.MyUDF';

select row, my_udf(row) from table;

我知道如何在 MR 引擎中执行 Hive UDF 中的 taskId：How to get the taskID or mapperID(something like partitionID in Spark) in a hive UDF?，但在 Spark 引擎中执行时无效。请告诉我如何在 Hive UDF 中获取 Spark 的 partitionID 或 taskContext，非常感谢！

【问题讨论】：

UDF 无权访问 Spark 的内部。如果您的数据在数据集中，您可以使用 spark_partition_id 函数。
@kanielc 你的意思是没有办法在 Hive UDF 中获取 Spark taskContext ？
我自己有正确答案，看这里：stackoverflow.com/questions/68079180/…

标签： apache-spark hive apache-spark-sql user-defined-functions

【解决方案1】：

这也许可以做你想做的事（下面的示例 UDF）

import org.apache.spark.TaskContext

val foo = udf  { (x: Int) => TaskContext.getPartitionId }

getPartitionId 函数应该会得到你需要的信息。

【讨论】：

感谢您的回答，但这不是我需要的。我想在 Hive UDF 中获取 Spark partitionID 或 taskContext。
stackoverflow.com/questions/68079180/… ，我自己在这里找到了答案。

【解决方案2】：

我自己有正确答案，请看这里：How to get the taskID or mapperID(something like partitionID in Spark) in a hive UDF?

上面的这个答案将为您介绍以下提示：

如何使用 MR 执行引擎在 Hive UDF 中获取 taskId；
如何使用 Spark 执行引擎在 Hive UDF 中获取 taskContext 和 partitionId。

希望它对你有用！

【讨论】：