[Python] Python/Pyspark, return geojson record with Z values / Truncate Z

I have a dataframe with a column called 'geometry', it contains Multipolygon and Polygon values. I want to use Pyspark or Python to find and return where coordinates contain [X, Y, Z]; also would like to create a code block to remove the Z value if exists. How can I do this? Example below, I'd like to return the first coordinate value, and not return any other values.

I want to do something like below but can't figure out how to append a new column to the dataframe, the coding required to look for and return only rows with X,Y,Z geometry:

Code I'm thinking of using to find Z value, but does not work:

for row in source_df.collect():
z_val = len(row.geometry["coordinates"]) =3 for x in row.geometry["coordinates"]]


Example data:

{
"type": "MultiPolygon",
"coordinates": [
[-120.92484404138442,35.54577502278743,0.0],
[-120.92484170835023,35.545764670080004],
[-120.92470946198651,35.54517811398435],
[-120.92373579577058,35.54476080459215],
[-120.92224560209857,35.544644824151],
[-120.91471743922112,35.54405891151482],
[-120.9137131887035,35.541405607829184],
[-120.91370267246779,35.54138005556737],
[-120.91368022915093,35.54133577314701],
[-120.91365314934913,35.54129325687539],
[-120.91364620938849,35.541283659095036],
[-120.91019544280519,35.53661949063082],
[-120.91016692865233,35.536584105321104],
[-120.91013516362523,35.53655061941634],
[-120.9101046793985,35.53652289241281],
[-120.90545581970368,35.53257237955164],
[-120.90540343303125,35.53253236763702]
]
}

Continue reading...