title: pyspark group and split dataframe
body: <p>I am trying to filter and then split a dataset into two separate files.</p>
  
  <p>Dataset: test.txt (Schema: uid, prod, score)</p>
  
  <pre><code>1   XYZ 2.0
  2   ABC 0.5
  1   PQR 1.0
  2   XYZ 2.1
  3   PQR 0.5
  1   ABC 0.5
  </code></pre>
  
  <p>First, I want to filter any uid having less than or equal to 1 product. I have already achieved that by the following code.</p>
  
  <pre><code>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 &gt;= 2")
  all_data = df.join(filtered, df.uid == filtered.uid , 'inner').drop(filtered.uid).drop(filtered.n)
  all_data.show()
  </code></pre>
  
  <p>This produces the following output:</p>
  
  <pre><code>+----+-----+---+
  |prod|score|uid|
  +----+-----+---+
  | XYZ|  2.0|  1|
  | PQR|  1.0|  1|
  | ABC|  0.5|  1|
  | ABC|  0.5|  2|
  | XYZ|  2.1|  2|
  +----+-----+---+
  </code></pre>
  
  <p>I need to now create 2 files out of the above dataframe. The problem that I am now facing is what is the best way to take one row for each product (can be any row) and put it in a different file(val.txt) and the rest of the rows in a different file (train.txt).</p>
  
  <p>Expected output (train.txt)</p>
  
  <pre><code>1    XYZ    2.0
  1    PQR    1.0
  2    ABC    0.5
  </code></pre>
  
  <p>Expected output (val.txt)</p>
  
  <pre><code>1    ABC    0.5
  2    XYZ    2.1
  </code></pre>
  
  <p>Thanks in advance !</p>