Applying a distinct operation within each rdd

user_2939259 0
Viewed 160

I have some very large pyspark dataframes with many duplicate rows. However, in my use case, the shuffle to do a full distinct() is not worth it in time cost. Instead, I wish to just apply distinct within each partition. I can't figure out how to do this, however. I've tried:

>>> spark = SparkSession.builder.appName('foobar').getOrCreate()
>>> data = [("Java", "20000"), ("Python", "100000"), ("Scala", "3000")]
>>> df = spark.createDataFrame(data)
>>> df.rdd.mapPartitions(lambda p: p.distinct()).collect()

and

>>> spark = SparkSession.builder.appName('foobar').getOrCreate()
>>> data = [("Java", "20000"), ("Python", "100000"), ("Scala", "3000")]
>>> df = spark.createDataFrame(data)
>>> df.foreachPartition(lambda p: p.distinct()).collect()

but in both cases I get

AttributeError: 'itertools.chain' object has no attribute 'distinct'

Would you please have advice on how to achieve this?

pyspark apache-spark-sql apache-spark python
1 Answers

The type of p in mapPartitions is an itertools.chain, not an RDD. You can return the unique elements from a chain with a function like below and then use it in the call to mapPartitions:

def uniq(chain):
    seen = set()
    for x in chain:
        if x in seen:
            continue
        seen.add(x)
        yield x

rdd = sc.parallelize([1,2,2,3,3,4])
rdd.repartition(1).mapPartitions(uniq).foreach(print)

# outputs
# 1
# 2
# 3
# 4

I've repartitioned into a single partition above to simulate uniqueness across the whole example data set but this will work correctly assuming you have already partitioned your data appropriately.

EDIT:

Performance

I don't know spark perf that well but in general the uniq function above should be faster than distinct because we are explicitly operating on a single partition rather than spark having to coordinate the distinctness across multiple partitions.

Even if there only happens to be 1 partition I suspect there is still a layer of coordination as indicated by the below results - I'm actually pretty surprised how much faster using uniq was over distinct on a single partition...

The below isn't concrete due to the very low sample size running locally but hopefully it gives a rough indicator.

# setup
rdd = sc.parallelize(["hello"] * 1000000).repartition(1)
rdd2 = rdd.repartition(2)

import timeit

# uniq 1 partition
timeit.timeit(lambda: rdd.mapPartitions(uniq).foreach(print), number=100)
# snip output...
17.04100012999993

# uniq 2 partitions
timeit.timeit(lambda: rdd2.mapPartitions(uniq).foreach(print), number=100)
# snip output...
9.535805986000014

# distinct 1 partition
timeit.timeit(lambda: rdd.distinct().foreach(print), number=100)
# snip output...
74.313582924

# distinct 2 partitions
timeit.timeit(lambda: rdd2.distinct().foreach(print), number=100)
# snip output...
38.47051327800011

So the mapPartitions(uniq) approach seems to be a fair bit more performant (bearing in mind the above is a pretty limited test).

user_141204630
Related