Question: I would like to get the first and last row of each partition in spark (I'm using pyspark). How do I go about this?
In my code I repartition my dataset based on a key column using:

`mydf.repartition(keyColumn).sortWithinPartitions(sortKey)`

Is there a way to get the first row and last row for each partition?
Thanks

Answer: 
I would highly advise against working with partitions directly. Spark does a lot of DAG optimisation, so when you try executing specific functionality on each partition, all your assumptions about the partitions and their distribution might be completely false.

You seem to however have a `keyColumn` and `sortKey`, so then I'd just suggest to do the following:

```
import pyspark
import pyspark.sql.functions as f

w_asc = pyspark.sql.Window.partitionBy(keyColumn).orderBy(f.asc(sortKey))
w_desc = pyspark.sql.Window.partitionBy(keyColumn).orderBy(f.desc(sortKey))
res_df = mydf. \
 withColumn("rn_asc", f.row_number().over(w_asc)). \
 withColumn("rn_desc", f.row_number().over(w_desc)). \
 where("rn_asc = 1 or rn_desc = 1")

```

The resulting dataframe will have 2 additional columns, where `rn_asc=1` indicates the first row and `rn_desc=1` indicates the last row.
Here is another approach using [mapPartitions](https://stackoverflow.com/questions/26741714/how-does-the-pyspark-mappartitions-function-work) from RDD API. We iterate over the elements of each partition until we reach the end. I would expect this iteration to be very fast since we skip all the elements of the partition except the two edges. Here is the code:

```py
df = spark.createDataFrame([
  ["Tom", "a"],
  ["Dick", "b"],
  ["Harry", "c"],
  ["Elvis", "d"],
  ["Elton", "e"],
  ["Sandra", "f"]
], ["name", "toy"])

def get_first_last(it):
      first = last = next(it)
      for last in it:
        pass

      # Attention: if first equals last by reference return only one!
      if first is last:
        return [first]

      return [first, last]

# coalesce here is just for demonstration
first_last_rdd = df.coalesce(2).rdd.mapPartitions(get_first_last)

spark.createDataFrame(first_last_rdd, ["name", "toy"]).show()

# +------+---+
# |  name|toy|
# +------+---+
# |   Tom|  a|
# | Harry|  c|
# | Elvis|  d|
# |Sandra|  f|
# +------+---+

```

PS: Odd positions will contain the first partition element and the even ones the last item. Also note that the number of results will be `(numPartitions * 2) - numPartitionsWithOneItem` which I expect to be relatively small therefore you shouldn't bother about the cost of the new `createDataFrame` statement.