Perform multiple column lookups based on ID

I have a pyspark dataframe with several columns:

+-------------+----------+------+
|.        id  |      num |cat   |
+-------------+----------+------+
|        00111|     50012|     a|
|        00111|     10131|     a|
|        00111|     11001|     b| 
|        10131|     71010|     a|
|        10131|     60010|     c|
|        11001|     53420|     z|
|        11001|     20011|     a|
|        11001|     00000|     q|
|        13403|     33001|     a|
|        13403|     10023|     a|
|        50012|     00111|     a|
+-------------+----------+------+

I would like to do the following:

1. for each unique entry in id select ALL values in num
2. look up the rows in id that match the selected values in num
3. select the rows in cat that correspond to the num values selected and count the number of occurrences of different values.

For example:

1. id==00111 and num = [50012; 10131; 11001]

2. select rows where id == num

   +-------------+----------+------+
   |.        id  |      num |cat   |
   +-------------+----------+------+
   |        10131|     71010|     a|
   |        10131|     60010|     c|
   |        11001|     53420|     z|
   |        11001|     20011|     a|
   |        11001|     00000|     q|
   |        50012|     00111|     a|
   +-------------+----------+------+
   

3. select the number of times cat = "a" which would be 3 in this case.

4. repeat for each unique value in id

The output would look something like this:

+-------------+----------+------+------+
|.        id  |      num |cat   |cat   |
+-------------+----------+------+------+
|        00111|     50012|     a|     3|
|        00111|     10131|     a|     3|
|        00111|     11001|     b|     3|
|        10131|     71010|     a|     1|
|        10131|     60010|     c|     1|
|        11001|     53420|     z|     0|
|        11001|     20011|     a|     0|
|        11001|     00000|     q|     0|
|        13403|     33001|     a|     0|
|        13403|     10023|     a|     0|
|        50012|     00111|     a|     2|
+-------------+----------+------+------+

Try this.

First, map the num to id with left join, then count when only cat is "a" which has the same id.

w = Window.partitionBy('id')
df = (df.join(df.select(*[F.col(x).alias(f'{x}_right') for x in df.columns]), 
              on=F.col('num') == F.col('id_right'), how='left')
      .select(*df.columns,
              F.count(F.when(F.col('cat_right') == 'a', 1)).over(w).alias('cnt'))
      .dropDuplicates())

Result

+-----+-----+---+---+
|   id|  num|cat|cnt|
+-----+-----+---+---+
|00111|50012|  a|  3|
|00111|10131|  a|  3|
|00111|11001|  b|  3|
|10131|71010|  a|  0|
|10131|60010|  c|  0|
|11001|53420|  z|  0|
|11001|20011|  a|  0|
|11001|00000|  q|  0|
|13403|33001|  a|  0|
|13403|10023|  a|  0|
|50012|00111|  a|  2|
+-----+-----+---+---+

Join id and num by self, and just count the cat.

w = Window.partitionBy('id')
df1 = df.drop('cat')
df2 = df.drop('id').withColumnRenamed('num', 'id')
df3 = df1.join(df2, ['id'], 'inner')
df3.show()
df3.withColumn('cnt', f.count(f.when(f.col('cat') == f.lit('a'), True)).over(w)) \
.show()
+-----+-----+---+
|   id|  num|cat|
+-----+-----+---+
|00111|50012|  p|
|00111|10131|  p|
|00111|11001|  p|
|10131|71010|  a|
|10131|60010|  a|
|11001|53420|  b|
|11001|20011|  b|
|11001|00000|  b|
|50012|00111|  a|
+-----+-----+---+
+-----+-----+---+---+
|   id|  num|cat|cnt|
+-----+-----+---+---+
|00111|50012|  p|  0|
|00111|10131|  p|  0|
|00111|11001|  p|  0|
|10131|71010|  a|  2|
|10131|60010|  a|  2|
|11001|53420|  b|  0|
|11001|20011|  b|  0|
|11001|00000|  b|  0|
|50012|00111|  a|  1|
+-----+-----+---+---+