from decimal import Decimal
from pydantic import Field
from pyspark.sql import SparkSession
from pyspark.sql.types import StructField, IntegerType, StringType, DecimalType, StructType
from pyspark.testing import assertDataFrameEqual
from sparkdantic import SparkModel
class FooBar(SparkModel):
id: int
name: str
threshold: Decimal = Field(decimal_places=2)
def test_decimal():
spark = SparkSession.builder.master('local[*]').getOrCreate()
expected_schema = StructType([StructField('id', IntegerType(), False),
StructField('name', StringType(), False),
StructField('threshold', DecimalType(10, 2), False)])
foobar = FooBar(id=95, name='test name', threshold=Decimal('0.33'))
df_expected = spark.createDataFrame([(95, 'test name', Decimal('0.33'))],
expected_schema)
df = spark.createDataFrame([foobar.model_dump()], FooBar.model_spark_schema())
df.show(truncate=False)
df_expected.show(truncate=False)
print(df.schema)
print(df_expected.schema)
assertDataFrameEqual(df, df_expected)
spark.stop()
+---+---------+---------+
|id |name |threshold|
+---+---------+---------+
|95 |test name|0 |
+---+---------+---------+
+---+---------+---------+
|id |name |threshold|
+---+---------+---------+
|95 |test name|0.33 |
+---+---------+---------+
StructType([StructField('id', IntegerType(), False), StructField('name', StringType(), False), StructField('threshold', DecimalType(10,0), False)])
StructType([StructField('id', IntegerType(), False), StructField('name', StringType(), False), StructField('threshold', DecimalType(10,2), False)])
E pyspark.errors.exceptions.base.PySparkAssertionError: [DIFFERENT_ROWS] Results do not match: ( 100.00000 % )
E *** actual ***
E ! Row(id=95, name='test name', threshold=Decimal('0'))
E
E
E *** expected ***
E ! Row(id=95, name='test name', threshold=Decimal('0.33'))
I'm not sure if it's me doing it wrong or there is a bug here.