PySpark - RDD contains an integer column pointing to float values in memory, cannot create DataFrame

Viewed 253

I have encountered a bizarre error when trying to use an RDD to create a PySpark DataFrame. Normally spark.createDataFrame(df.rdd, new_schema) works fine as long as the schema is compatible with the RDD. In the following case, though, the RDD has values that should be integers but are instead floats, which causes an error.

I believe that the RDD is trying to do some tricky memory optimization by only storing unique values once (per partition? per block?) and having each cell in the DataFrame point to the same address in memory. It seems to consider 1 and 1.0 to be the same value despite having different data types.

I would like to fundamentally understand WHY this is happening, and whether it is a bug. But more pertinently, how can I work around it? Can I prevent the RDD from doing this in the first place? Can I somehow access df.rdd without running into this? Can I "typecast" the columns in the RDD before passing it to the dataframe constructor?

Reproducible case:

from pyspark.sql import types as T
from pyspark.sql import functions as F

# Create a simple dataframe with an integer and a float column
# This has the important traits:
#   a) both columns contain the same value: 1 or 1.0
#   b) that value is repeated twice in the first column, which seems to trigger spark
#      to optimize memory by pointing to the same address in memory
df = spark.createDataFrame(
    data=[
        (1, 3.1415),
        (2, 1.0000),
        (1, 3.1415),
    ],
    schema=T.StructType(
        [
            T.StructField("x", T.IntegerType(), True),
            T.StructField("y", T.DoubleType(), True),
        ]
    ),
)
df = df.repartition(1) # this is needed to get all the data into one block(?)
df.show()

# Build a second dataframe with same data, but nested within structs
df2 = df.withColumn("x2", F.struct(F.col("x")))
df2 = df2.withColumn("y2", F.struct(F.col("y")))
df2 = df2.select("x2", "y2")
df2.show()

# Print the rows in each dataframe. The second will have row 3, column 1 pointing to
# the same value in memory as row 2, column 2.
print("df1:")
for row in df.collect():
    print(row)

print("\ndf2:")
for row in df2.collect():
    print(row)

# Try to build a new dataframe from the RDD. We will get an error
spark.createDataFrame(df2.rdd, df2.schema).show()

This prints:

df1
Row(x=1, y=3.1415)
Row(x=2, y=1.0)
Row(x=1, y=3.1415)

df2
Row(x2=Row(x=1), y2=Row(y=3.1415))
Row(x2=Row(x=2), y2=Row(y=1.0))
Row(x2=Row(y=1.0), y2=Row(y=3.1415))

Note that in df2, which has the values nested under structs, the third value for x2 is 1.0, even though that column has IntegerType! (Printing the memory addresses confirms that it's the same pointer as y2 row 2). This results in the following error when the RDD is used:

TypeError: field x in field x2: IntegerType can not accept object 1.0 in type <class 'float'>
0 Answers
Related