Нужен список столбцов, значение которых больше 0 в pyspark

Question

У меня есть данные ниже:

>>> dfStd1.show()
+---+----+------+-------+-----------------------------------------------+------+
| id|Name|Seq_Id|Carrier|CASE WHEN (NOT (Seq_Id = 1)) THEN 0 ELSE 12 END|string|
+---+----+------+-------+-----------------------------------------------+------+
|  0|   0|     0|      2|                                              0|     0|
+---+----+------+-------+-----------------------------------------------+------+  

Итак, здесь мне нужны имена столбцов, значение которых больше 0. Здесь, например, для меня нужен столбец Carrier, мне нужно хранить такие значения всписок. Я попробовал приведенный ниже код, но не работает, а также сослался на множество ссылок SO, но не повезло:

>>> dfStd1[(dfStd1 > 0).any(axis=1)]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: '>' not supported between instances of 'DataFrame' and 'int'  

Выдает ошибку, как указано выше. Я даже пытался преобразовать его в панд, а затем отфильтровать, но безрезультатно.

Answer 1

Сначала вам нужны числовые столбцы:

schema = {col: col_type for col, col_type in df.dtypes}
numeric_cols = [
            col
            for col, col_type in schema.items()
            if col_type in "int double bigint".split()
        ]

Затем вы можете подсчитать количество элементов в столбце, которые больше 0, используя:

from pyspark.sql.functions import when, col

count_cols_gt_zero = [
            json.loads(x)
            for x in self.data.select(
                [count(when(col(c) > 0, c)).alias(c) for c in self.schema]
            )
            .toJSON()
            .collect()
        ][0]

Затем, наконец:

final = [x for x, y in count_cols_gt_zero.items() if y > 0]