[Python] Hierarchy Transformation/Looping using Pyspark, Python

I have some SAS code and need to convert it to PySpark code.

Hierarchical SAS Code Part (1)

if cc1 = 1 then do i = 2, 3; cc(i) = 0; end;
if cc2 = 1 then do i = 3; cc(i) = 0; end;



Basically, this code sets cc(2) and cc(3) to 0 if cc1 (condition 1) is equal to 1. Then, the next line sets cc(3) to 0 if cc2 (condition 2) is equal to 1, and so forth. In this chunk of code, there is a clear hierarchical order [cc1 > cc2 > cc3].

Hierarchical PySpark Code Part (1) Note: Use AI to summarize the purpose of my code if needed.

To replicate Hierarchical SAS Code Part (1), I wrote the following PySpark script:


data = [('A', 1, 0, 1, 1, 0, 1),
('B', 0, 1, 0, 1, 0, 0) ]
columns_t = ['id', "c1", "c2", "c3", "c4", 'c5', 'c6']

test_data = spark.createDataFrame(data, columns_t)

array_1 = ["c1", "c2", "c3"]
array_2 = ['c4', "c5", "c6"]

array_dict = {"array_1": array_1, "array_2": array_2}

for key in array_dict.keys():
if key not in test_data.columns:
test_data = test_data.withColumn(key, lit(None))

for key, array in array_dict.items():
for x in array:
test_data = test_data.withColumn(
key,
when(col(key).isNull() & (col(x) == 1), lit(x)).otherwise(col(key))
)

selected_columns = ["id"] + list(array_dict.keys())

from pyspark.sql.functions import when, col, lit

new_data_dict = test_data.select(
*selected_columns,
*[
when(col(key) == f'{x}', lit(1)).otherwise(lit(0)).alias(x)
for key, array in array_dict.items() for x in array
]
)

display(test_data)
display(new_data_dict)


However, Hierarchical PySpark Code Part (1) does not work for Hierarchical SAS Code Part (2), as seen below, because there is no clear form of hierarchy.

Hierarchical SAS Code Part (2)

if cc397 = 1 then do i = 303,398, 399; cc(i) = 0; end;
if cc398 = 1 then do i = 303,399; cc(i) = 0; end;


In the SAS code above, if cc397 equals 1, then set cc303, cc398, and cc399 to 0. If cc398 equals 1, then set cc303 and cc399 to 0.

Hierarchical Order in SAS Code Part (2): [cc397 > cc398 > cc303, cc399]

Ultimately, I need some help developing PySpark code to take into account the SAS Code Part 2 scenario, or maybe some advice on making PySpark Code Part 1 more efficient.

P.S. I'm working with rather large data tables (50GB).

Please code above to see what I have done.

Continue reading...