如何在两行中将pyspark数据帧切片

Question

我在Databricks工作。

我有一个包含500行的数据框，我想创建两个包含100行的数据框，另一个包含剩余的400行。

+--------------------+----------+
|              userid| eventdate|
+--------------------+----------+
|00518b128fc9459d9...|2017-10-09|
|00976c0b7f2c4c2ca...|2017-12-16|
|00a60fb81aa74f35a...|2017-12-04|
|00f9f7234e2c4bf78...|2017-05-09|
|0146fe6ad7a243c3b...|2017-11-21|
|016567f169c145ddb...|2017-10-16|
|01ccd278777946cb8...|2017-07-05|

我尝试了以下操作，但收到错误消息

df1 = df[:99]
df2 = df[100:499]


TypeError: unexpected item type: <type 'slice'>

Answer 1

最初，我误会了，并以为您想分割这些列。如果要选择行的子集，一种方法是使用创建索引列monotonically_increasing_id()。从文档：

保证生成的ID是单调递增且唯一的，但不是连续的。

您可以使用此ID对数据框进行排序，并使用该ID对其子集进行排序，limit()以确保准确获得所需的行。

例如：

import pyspark.sql.functions as f
import string

# create a dummy df with 500 rows and 2 columns
N = 500
numbers = [i%26 for i in range(N)]
letters = [string.ascii_uppercase[n] for n in numbers]

df = sqlCtx.createDataFrame(
    zip(numbers, letters),
    ('numbers', 'letters')
)

# add an index column
df = df.withColumn('index', f.monotonically_increasing_id())

# sort ascending and take first 100 rows for df1
df1 = df.sort('index').limit(100)

# sort descending and take 400 rows for df2
df2 = df.sort('index', ascending=False).limit(400)

只是为了验证这是否符合您的要求：

df1.count()
#100
df2.count()
#400

我们还可以验证索引列是否不重叠：

df1.select(f.min('index').alias('min'), f.max('index').alias('max')).show()
#+---+---+
#|min|max|
#+---+---+
#|  0| 99|
#+---+---+

df2.select(f.min('index').alias('min'), f.max('index').alias('max')).show()
#+---+----------+
#|min|       max|
#+---+----------+
#|100|8589934841|
#+---+----------+