这是一种使用“colllect\u set”函数和一些pythonic操作的简单方法：

idLimit=3 #define your limit

id_lst=(sourceDF  #collect a list of distinct ids
        .select(collect_set('id'))
        .collect()[0][0]
       )

id_lst.sort() #sort the ids alphabatically

id_lst_limited=id_lst[:idLimit] #limit the list as per your defined limit

targetDF=(sourceDF #filter the source df using your limited list
          .filter("id in ({0})".format(str(id_lst_limited)[1:-1]))
         )

您可以使用sparksql查询来实现这一点。
换个衣服就行了 limit 子查询中的子句值，以选择不同的 id .

df=spark.createDataFrame([("A", "x", "3"), ("A", "y", "1"), ("B", "a", "2"), ("B", "b", "5"), ("C", "v", "2"), ("D", "f", "6")], ["id", "manager", "score"])

+---+-------+-----+
| id|manager|score|
+---+-------+-----+
|  A|      x|    3|
|  A|      y|    1|
|  B|      a|    2|
|  B|      b|    5|
|  C|      v|    2|
|  D|      f|    6|
+---+-------+-----+

df.createOrReplaceTempView("employee")

sql("""select * from employee e1 
     inner join (  
        select distinct id as uni_id from employee order by uni_id limit 3) e2 
     on e1.id=e2.uni_id""").drop("uni_id").show() 

+---+-------+-----+
| id|manager|score|
+---+-------+-----+
|  A|      x|    3|
|  A|      y|    1|
|  B|      a|    2|
|  B|      b|    5|
|  C|      v|    2|
+---+-------+-----+

你也可以试试 StringIndexer 将每个字符串分解为 id 列，然后 filter 根据限制。

import pyspark.sql.functions as F
from pyspark.ml.feature import StringIndexer

n = 3 #change as per limit
idx = StringIndexer(inputCol="id",outputCol="id_num")
idx.fit(df).transform(df).filter(F.col("id_num")<n).drop("id_num").show()

+---+-------+-----+
| id|manager|score|
+---+-------+-----+
|  A|      x|    3|
|  A|      y|    1|
|  B|      a|    2|
|  B|      b|    5|
|  C|      f|    2|
+---+-------+-----+

我注意到上面的一个答案是基于sparksql的。
下面是另一种基于sparksql的方法，但是 WINDOW 条款-

sql("select id, manager, score from (select e1.id, e1.manager, e1.score, dense_rank() over (order by e1.id) as rrank from employee e1) where rrank <= 3").show()

完整代码-

df=spark.createDataFrame([("A", "x", "3"), ("A", "y", "1"), ("B", "a", "2"), ("B", "b", "5"), ("C", "v", "2"), ("D", "f", "6")], ["id", "manager", "score"])

+---+-------+-----+
| id|manager|score|
+---+-------+-----+
|  A|      x|    3|
|  A|      y|    1|
|  B|      a|    2|
|  B|      b|    5|
|  C|      v|    2|
|  D|      f|    6|
+---+-------+-----+

df.createOrReplaceTempView("employee")

sql("select id, manager, score from (select e1.id, e1.manager, e1.score, dense_rank() over (order by e1.id) as rrank from employee e1) where rrank <= 3").show()

+---+-------+-----+
| id|manager|score|
+---+-------+-----+
|  A|      x|    3|
|  A|      y|    1|
|  B|      a|    2|
|  B|      b|    5|
|  C|      v|    2|
+---+-------+-----+

你可以简单地避免 join 由 where id in (select distinct id ...limit 3) 如下所示-

val df = Seq(("A", "x", "3"), ("A", "y", "1"), ("B", "a", "2"), ("B", "b", "5"), ("C", "v", "2"), ("D", "f", "6"))
    .toDF("id", "manager", "score")
    df.show(false)

    /**
      * +---+-------+-----+
      * |id |manager|score|
      * +---+-------+-----+
      * |A  |x      |3    |
      * |A  |y      |1    |
      * |B  |a      |2    |
      * |B  |b      |5    |
      * |C  |v      |2    |
      * |D  |f      |6    |
      * +---+-------+-----+
      */

    df.createOrReplaceTempView("employee")
    spark.sql("select * from employee where id in (select distinct id from employee order by id limit 3)")
      .show(false)

    /**
      * +---+-------+-----+
      * |id |manager|score|
      * +---+-------+-----+
      * |A  |x      |3    |
      * |A  |y      |1    |
      * |B  |a      |2    |
      * |B  |b      |5    |
      * |C  |v      |2    |
      * +---+-------+-----+
      */