python - 如何获取 row_number 是 pyspark 数据框

Question

为了排名，我需要获取 row_number 是一个 pyspark 数据框。我看到 pyspark 的 windows 函数中有 row_number 函数，但这需要使用 HiveContext。

我尝试用 HiveContext 替换 sqlContext

        import pyspark
        self.sc = pyspark.SparkContext()
        #self.sqlContext = pyspark.sql.SQLContext(self.sc)
        self.sqlContext = pyspark.sql.HiveContext(self.sc)

但它现在抛出异常 TypeError: 'JavaPackage' object is not callable您能否帮助操作 HiveContext 或以不同的方式获取行号？

数据示例:我想先根据我的预测进行排名，然后根据这个排名计算损失函数 (ndcg)。为了计算损失函数，我需要排名(即预测在排序中的位置)

所以第一步是按 pred 对数据进行排序，然后我需要一个排序数据的运行计数器。

+-----+--------------------+
|label|pred|
+-----+--------------------+

|  1.0|[0.25313606997906...|
|  0.0|[0.40893413256608...|
|  0.0|[0.18353492079000...|
|  0.0|[0.77719741215204...|
|  1.0|[0.62766290642569...|
|  1.0|[0.40893413256608...|
|  1.0|[0.63084085591913...|
|  0.0|[0.77719741215204...|
|  1.0|[0.36752166787523...|
|  0.0|[0.40893413256608...|
|  1.0|[0.25528507573737...|
|  1.0|[0.25313606997906...|

谢谢。

Answer 1

您不需要创建 HiveContext如果您的数据不在 Hive 中。您可以继续使用 sqlContext。

除非您创建数据框，否则没有row_number。 pyspark.sql.functions.row_number ` 用于不同的目的，它仅适用于窗口分区。

您可能需要使用 monotonically_increasing_id 创建一个新列作为 row_id稍后再查询。

from pyspark.sql.functions import monotonically_increasing_id
from pyspark.sql.types import Row

data = sc.parallelize([
  Row(key=1, val='a'),
  Row(key=2, val='b'), 
  Row(key=3, val='c'), 
]).toDF()

data = data.withColumn(
  'row_id',
  monotonically_increasing_id()
)

data.collect()


Out[8]: 
[Row(key=1, val=u'a', row_id=17179869184),
 Row(key=2, val=u'b', row_id=42949672960),
 Row(key=3, val=u'c', row_id=60129542144)]