Coverage for jstark/features/count.py: 100%

14 statements  

« prev     ^ index     » next       coverage.py v7.4.3, created at 2024-02-25 20:09 +0000

1"""Count feature""" 

2from typing import Callable 

3import pyspark.sql.functions as f 

4from pyspark.sql import Column 

5 

6 

7from .feature import BaseFeature 

8 

9 

10class Count(BaseFeature): 

11 def aggregator(self) -> Callable[[Column], Column]: 

12 return self.count_aggregator 

13 

14 def column_expression(self) -> Column: 

15 return f.lit(1) 

16 

17 def default_value(self) -> Column: 

18 return f.lit(0) 

19 

20 @property 

21 def description_subject(self) -> str: 

22 return "Count of rows"