Подсчет нулей в фреймах данных PySpark с общим количеством строк и столбцов

Question

Я пытаюсь написать запрос для подсчета всех значений null в большом фрейме данных, используя PySpark. После прочтения в наборе данных я делаю это:

import pyspark.sql.functions as F

df_agg = df.agg(*[F.count(F.when(F.isnull(c), c)).alias(c) for c in df.columns])

df_countnull_agg.coalesce(1).write.option("header", "true").mode("overwrite").csv(path)

Это прекрасно работает, и df_agg dataframe дает мне что-то вроде этого:

#+--------+--------+--------+
#|Column_1|Column_2|Column_3|
#+--------+--------+--------+
#|      15|      56|      18|
#+--------+--------+--------+

Что я хочу сделать, это также добавить два столбца в конце кадра данных для total_rows и total_columns, чтобы я мог выполнить некоторые вычисления после записи в файл .csv. Я знаю, что могу получить числа из фрейма данных следующим образом:

total_rows = df.count()
total_columns = len(df.columns)

Я хочу добавить эти два числа в столбцы, которые приведут к такому фрейму данных, а затем записать его в .csv как Я до:

#+--------+--------+--------+--------+--------+
#|Column_1|Column_2|Column_3|t_rows  |t_cols  |
#+--------+--------+--------+--------+--------+
#|      15|      56|      18|     500|      20|
#+--------+--------+--------+--------+--------+

Что меня беспокоит, так это время выполнения, так как подсчет nulls занимает немного времени, а затем вычисляет форму фрейма данных и добавляет его в конечный df для вывода , Любая помощь приветствуется!

Answer 1

Чтобы подсчитать общее количество строк, вы можете сделать это внутри агрегата, посчитав значения F.lit(1), а затем вы можете получить count от общего числа столбцов, используя withColumn для создания нового столбца с литералом (lit) как len из df.columns.

df.agg(*[F.count(F.when(F.isnull(c), c)).alias(c) for c in df.columns], F.count(F.lit(1)).alias("t_rows"))\
   .withColumn("t_cols", F.lit(len(df.columns))).show()

+-----+----+--------+------+------+
|query|href|position|t_rows|t_cols|
+-----+----+--------+------+------+
|    3|   2|       0|    12|     3|
+-----+----+--------+------+------+