如何获取 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

如果您的数据不在 Hive 中，则无需创建 HiveContext。你可以继续使用你的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)]