pyspark組和拆分數據框

Question

我正在嘗試過濾，然后將數據集分成兩個單獨的文件。

數據集：test.txt（架構：uid，prod，score）

1   XYZ 2.0
2   ABC 0.5
1   PQR 1.0
2   XYZ 2.1
3   PQR 0.5
1   ABC 0.5

首先，我要過濾任何小於或等於1個乘積的uid。 我已經通過以下代碼實現了這一點。

from pyspark.sql.types import *
from pyspark.sql.functions import *

rdd = sc.textFile('test.txt').map(lambda row: row.split('\t'))
schema = StructType([
           StructField('uid', IntegerType(), True),
           StructField('prod', StringType(), True),
           StructField('score', FloatType(), True)])
df = rdd.toDF([f.name for f in schema.fields])
filtered = df.groupby('uid').count().withColumnRenamed("count", "n").filter("n >= 2")
all_data = df.join(filtered, df.uid == filtered.uid , 'inner').drop(filtered.uid).drop(filtered.n)
all_data.show()

這將產生以下輸出：

+----+-----+---+
|prod|score|uid|
+----+-----+---+
| XYZ|  2.0|  1|
| PQR|  1.0|  1|
| ABC|  0.5|  1|
| ABC|  0.5|  2|
| XYZ|  2.1|  2|
+----+-----+---+

我現在需要從上述數據框中創建2個文件。 我現在面臨的問題是，對每種產品排一行（可以是任何行）並將其放在不同的文件（val.txt）中，而將其余的行放在不同的文件（火車中）的最佳方法是什么。文本）。

預期輸出（train.txt）

1    XYZ    2.0
1    PQR    1.0
2    ABC    0.5

預期輸出（val.txt）

1    ABC    0.5
2    XYZ    2.1

提前致謝 ！

Answer 1

我認為這里的關鍵問題是您沒有數據的主鍵。

all_data = all_data.withColumn(
    'id',
    monotonically_increasing_id()
)

train = all_data.dropDuplicates(['prod'])

# could OOM if your dataset is too big
# consider BloomFilter if so
all_id = {row['id'] for row in all_data.select('id').collect()}
train_id = {row['id'] for row in train.select('id').collect()}
val_id = all_id - train_id

val = all_data.where(col(id).isin(val_id))