Unittest (PyTest) function used in mapPartitions() in PySpark - pyspark

Question

I have a use case where in I'm writing a dataframe to DynamoDB.
The issue is, that although I'm able to run the code on the cluster, the function used inside of mapPartitions() is getting hard to test.
Here is my code
import boto3
import json
def write_df(table_name, region):
the_list=[]
def my_map_partition(partition):
table = boto3.resource("dynamodb", region=region).Table(table_name)
for row in partition:
table.put_item(Item = json.loads(row.col1))
the_list.append(row.col1)
return the_list
return my_map_partition
def load_to_db(df):
df.rdd.mapPartitions(write_df("actual_table", "eu-west-2"))
#mock_dynamodb2
def init_dynamo():
client = boto3.client("dynamodb", "eu-west-2")
client.create_table(
TableName="my_test_table",
KeySchema=[
{"AttributeName": "partitionKey", "KeyType": "HASH"},
{"AttributeName": "sortKey", "KeyType": "RANGE"}
],
AttributeDefinitions=[
{"AttributeName": "partitionKey", "AttributeType": "S"},
{"AttributeName": "sortKey", "AttributeType": "S"},
]
)
def test_write_df():
self.init_dynamo()
columns = ['col1', 'col2']
vals = [('{"partitionKey":"pk1", "sortKey":"sk1", "random_value":"hello_world"}', 'dummy'),('{"partitionKey":"pk2", "sortKey":"sk2", "random_value":"hello_world"}', 'dummy')]
data = spark.createDataframe(vals, columns).select("col1")
a = data.rdd.mapPartitions(write_df("my_test_table", "eu-west-2"))
a.count()
pass #here is where I will assert values
Here is the error
objc[99130]: +[__NSCFConstantString initialize] may have been in progress in another thread when fork() was called.
objc[99130]: +[__NSCFConstantString initialize] may have been in progress in another thread when fork() was called. We cannot safely call it or ignore it in the fork() child process. Crashing instead. Set a breakpoint on objc_initializeAfterForkError to debug.
22/07/15 14:24:38 ERROR Executor: Exception in task 0.0 in stage 3.0 (TID 32)
org.apache.spark.SparkException: Python worker exited unexpectedly (crashed)
at org.apache.spark.api.python.BasePythonRunner$ReaderIterator$$anonfun$1.applyOrElse(PythonRunner.scala:585)
at org.apache.spark.api.python.BasePythonRunner$ReaderIterator$$anonfun$1.applyOrElse(PythonRunner.scala:567)
.
.
ERROR TaskSetManager: Task 0 in stage 3.0 failed 1 times; aborting job
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/usr/lib/python3.9/site-packages/pyspark/sql/session.py:66: in toDF
return sparkSession.createDataFrame(self, schema, sampleRatio)
/usr/lib/python3.9/site-packages/pyspark/sql/session.py:675: in createDataFrame
return self._create_dataframe(data, schema, samplingRatio, verifySchema)
/usr/lib/python3.9/site-packages/pyspark/sql/session.py:698: in _create_dataframe
rdd, schema = self._createFromRDD(data.map(prepare), schema, samplingRatio)
/usr/lib/python3.9/site-packages/pyspark/sql/session.py:486: in _createFromRDD
struct = self._inferSchema(rdd, samplingRatio, names=schema)
/usr/lib/python3.9/site-packages/pyspark/sql/session.py:460: in _inferSchema
first = rdd.first()
/usr/lib/python3.9/site-packages/pyspark/rdd.py:1588: in first
rs = self.take(1)
/usr/lib/python3.9/site-packages/pyspark/rdd.py:1568: in take
res = self.context.runJob(self, takeUpToNumLeft, p)
/usr/lib/python3.9/site-packages/pyspark/context.py:1227: in runJob
sock_info = self._jvm.PythonRDD.runJob(self._jsc.sc(), mappedRDD._jrdd, partitions)
/usr/lib/python3.9/site-packages/py4j/java_gateway.py:1309: in __call__
return_value = get_return_value(
/usr/lib/python3.9/site-packages/pyspark/sql/utils.py:111: in deco
return f(*a, **kw)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
answer = 'xro77'
gateway_client = <py4j.clientserver.JavaClient object at 0x124798820>
target_id = 'z:org.apache.spark.api.python.PythonRDD', name = 'runJob'
def get_return_value(answer, gateway_client, target_id=None, name=None):
"""Converts an answer received from the Java gateway into a Python object.
For example, string representation of integers are converted to Python
integer, string representation of objects are converted to JavaObject
instances, etc.
:param answer: the string returned by the Java gateway
:param gateway_client: the gateway client used to communicate with the Java
Gateway. Only necessary if the answer is a reference (e.g., object,
list, map)
:param target_id: the name of the object from which the answer comes from
(e.g., *object1* in `object1.hello()`). Optional.
:param name: the name of the member from which the answer comes from
(e.g., *hello* in `object1.hello()`). Optional.
"""
if is_error(answer)[0]:
if len(answer) > 1:
type = answer[1]
value = OUTPUT_CONVERTER[type](answer[2:], gateway_client)
if answer[1] == REFERENCE_TYPE:
> raise Py4JJavaError(
"An error occurred while calling {0}{1}{2}.\n".
format(target_id, ".", name), value)
E py4j.protocol.Py4JJavaError: An error occurred while calling
I'm not sure what is wrong in the unit tests since this is a working code.
Can someone guide me in the right direction on how do I test a function in PySpark that is being used for mapPartitions() ?