Pyspark - Index from monotonically_increasing_id changes after list aggregation

Viewed 395

I'm creating an index using the monotonically_increasing_id() function in Pyspark 3.1.1.

I'm aware of the specific characteristics of that function, but they don't explain my issue.

After creating the index I do a simple aggregation applying the collect_list() function on the created index.

If I compare the results the index changes in certain cases, that is specifically on the upper end of the long-range when the input data is not too small.

Full example code:

import random
import string

from pyspark.sql import SparkSession
from pyspark.sql import functions as f
from pyspark.sql.types import StructType, StructField, StringType

spark = SparkSession.builder\
    .appName("test")\
    .master("local")\
    .config('spark.sql.shuffle.partitions', '8')\
    .getOrCreate()

# Create random input data of around length 100000:
input_data = []
ii = 0
while ii <= 100000:
    L = random.randint(1, 3)
    B = ''.join(random.choices(string.ascii_uppercase, k=5))
    for i in range(L):
        C = random.randint(1,100)
        input_data.append((B,))
        ii += 1

# Create Spark DataFrame:
input_rdd = sc.parallelize(tuple(input_data))

schema = StructType([StructField("B", StringType())])
dg = spark.createDataFrame(input_rdd, schema=schema)

# Create id and aggregate:
dg = dg.sort("B").withColumn("ID0", f.monotonically_increasing_id())
dg2 = dg.groupBy("B").agg(f.collect_list("ID0"))

Output:
dg.sort('B', ascending=False).show(10, truncate=False)
dg2.sort('B', ascending=False).show(5, truncate=False)

This of course creates different data with every run, but if the length is large enough (problem appears slightly at 10000, but not at 1000), it should appear everytime. Here's an example result:

+-----+-----------+
|B    |ID0        |
+-----+-----------+
|ZZZVB|60129554616|
|ZZZVB|60129554617|
|ZZZVB|60129554615|
|ZZZUH|60129554614|
|ZZZRW|60129554612|
|ZZZRW|60129554613|
|ZZZNH|60129554611|
|ZZZNH|60129554609|
|ZZZNH|60129554610|
|ZZZJH|60129554606|
+-----+-----------+
only showing top 10 rows

+-----+---------------------------------------+
|B    |collect_list(ID0)                      |
+-----+---------------------------------------+
|ZZZVB|[60129554742, 60129554743, 60129554744]|
|ZZZUH|[60129554741]                          |
|ZZZRW|[60129554739, 60129554740]             |
|ZZZNH|[60129554736, 60129554737, 60129554738]|
|ZZZJH|[60129554733, 60129554734, 60129554735]|
+-----+---------------------------------------+
only showing top 5 rows

The entry ZZZVB has the three IDs 60129554615, 60129554616, and 60129554617 before aggregation, but after aggregation the numbers have changed to 60129554742, 60129554743, 60129554744.

Why? I can't imagine this is supposed to happen. Isn't the result of monotonically_increasing_id() a simple long that keeps its value after having been created?

EDIT: As expected a workaround is to coalesce(1) the DataFrame before creating the id.

1 Answers

dg and df2 are two different dataframes, each with their own DAG. These DAGs are executed independently from each other when an action on one of the dataframes is called. So each time show() is called, the DAG of the respective dataframe is evaluated and during that evaluation, f.monotonically_increasing_id() is called.

To prevent f.monotonically_increasing_id() being called twice, you could add a cache after the withColumn transformation:

dg = dg.sort("B").withColumn("ID0", f.monotonically_increasing_id()).cache()

With the cache, the result of the first evaluation of f.monotonically_increasing_id() is cached and reused when evaluating the second dataframe.

Related