PySpark - Combine calculation and kwargs in 'select' in a 'for' loop - select

Question

I was told using withColumn within a for loop is not a good idea in terms of performance and I try to replace it with select, as suggested. But I don't think it's a good idea to write all columns (and yes, I'm a bit lazy^^); I already have a dict with the list of columns and I try to use it, but I get an error at the second loop.
My main issue is that I need at least one column for the calculation, but **kwargs argument come after the others, so I try to make it dynamic: remove the column from the list, pass the column first, then my calculation, finally the other columns (of course, the loop is over the column list).
The code runs in a palantir-foundry code repository (I have no other pySpark environment to test).
Here is a snippet of my code, I hope it's enough to make it reproductible:
input_schema = {
'col_1': T.StringType(),
'col_2': T.StringType(),
'col_3': T.StringType(),
'col_4': T.StringType(),
'col_4': T.StringType(),
}
#tranform_df(
Output("output_df"),
nsw=Input("My_Input")
)
def my_compute_function(nsw):
nsw = nsw \
.withColumn("errors", F.lit(None).cast("array<string>")) \
.withColumn("error_check", F.lit("Pouet").cast("string"))
for col_name, test_list in validation_list.items():
l_input_schema = dict(input_schema)
del l_input_schema[col_name]
for to_test in test_list:
nsw = nsw \
.select(
*col_name,
F.when(
(nsw.error_check.isNotNull()),
F.when(nsw.errors.isNull(), F.array(nsw.error_check))
.otherwise(F.array_union(nsw.errors, F.array(nsw.error_check))),
)
.otherwise(nsw.errors)
.alias("errors"),
**l_input_schema
)
return nsw
No syntax error here, but at the second loop, the following message:
Function select() got an unexpected argument col_2. Please review your code.
The error points out the line where kwargs is used (**l_input_schema), and it's always the second column

Answer 1

I feel a bit stupid because it was just a wrong way of using my dict as kwargs instead of args.
The following works fine:
nsw = nsw \
.select(
*input_schema,
F.when(
(nsw.error_check.isNotNull()),
F.when(nsw.errors.isNull(), F.array(nsw.error_check))
.otherwise(F.array_union(nsw.errors, F.array(nsw.error_check))),
)
.otherwise(nsw.errors)
.alias("errors"),
)
No need anymore to try to have the current column before calculation for each loop. On the other hand, all new columns need to be added before the loop, and the dict input_schema should also be updated accordingly.