Parallelize an operation applied to a list (PySpark)

At one point in my program, a function receives a list and repeats an operation using each of its items.

For the sake of this example, suppose my program starts off with a list l and counts the rows of some generic pyspark dataframe df that meet a condition.

l = [2, 4, 5]
res = []

for x in l:
    val = df.where((col('id') == x) | (col('id2') == x)).count()
    res.append(val)

Is it possible to have multiple workers calculate val simultaneously? That is, could each worker calculate its own val and append it to l individually?

This post suggests using foreach, but since I'm iterating over a list (not an RDD or dataframe), I cannot use that method.

you could do away with loops and use spark functions to count occurrences.

here's an example

# mock input data
data_sdf.show()

# +---+---+
# | id|id2|
# +---+---+
# |  2|  4|
# |  4|  1|
# |  3|  5|
# |  4|  5|
# |  1|  8|
# +---+---+

# create array of all ids, and remove ids that are not in that row
data_sdf. \
    withColumn('id_ls', func.array(*[func.lit(k) for k in ls])). \
    withColumn('id_per_row', func.expr('filter(id_ls, x -> (id = x) or (id2 = x))')). \
    filter(func.size('id_per_row') != 0). \
    select('*', func.explode('id_per_row').alias('ids')). \
    groupBy('ids'). \
    count(). \
    show()

# +---+-----+
# |ids|count|
# +---+-----+
# |4  |3    |
# |2  |1    |
# |5  |2    |
# +---+-----+

data_sdf. \
    withColumn('id_ls', func.array(*[func.lit(k) for k in ls])). \
    withColumn('id_per_row', func.expr('filter(id_ls, x -> (id = x) or (id2 = x))')). \
    show(truncate=False)

# +---+---+---------+----------+
# |id |id2|id_ls    |id_per_row|
# +---+---+---------+----------+
# |2  |4  |[2, 4, 5]|[2, 4]    |
# |4  |1  |[2, 4, 5]|[4]       |
# |3  |5  |[2, 4, 5]|[5]       |
# |4  |5  |[2, 4, 5]|[4, 5]    |
# |1  |8  |[2, 4, 5]|[]        |
# +---+---+---------+----------+

IIUC, you can try array_intersect to find the matching rows.

from pyspark.sql import functions as F

l = [2, 4, 5]
# Convert to Col type
l = [F.lit(x) for x in l]

df = (df.withColumn('intersect', F.array_intersect(F.array('id', 'id2'), F.array(l)))
      .select('id', F.explode('intersect').alias('intersect'))
      .groupby('intersect')
      .count())