Cdspt 15587 align functionality of totals methods for all extractor subclasses

mario/athena.py

@@ -1,5 +1,5 @@
 from pyathena import connect
-from mario.data_extractor import DataExtractor, Configuration
+from mario.data_extractor import DataExtractor, Configuration, GET_TOTAL_WITHOUT_MEASURE
 import logging
 import datetime
 import numbers
@@ -56,18 +56,11 @@ def get_connection(self):
             catalog_name=cfg.catalog
         )
 
-    def get_total(self, measure=None):
-        """
-        For totals when streaming data we need to run a totals SQL query separate
-        from the main query and use the results of this
-        :return: the total value of the query
-        NOTE this is a direct copy from StreamingDataExtractor
-        """
+    def __get_total__(self, measure):
         from pandas import read_sql
-        logger.info("Building totals query")
-        measure = self.__get_measure__(measure)
         if self.configuration.query_builder is not None:
             from mario.query_builder import QueryBuilder
+            logger.info("Building totals query")
             query_builder: QueryBuilder = self.configuration.query_builder(
                 configuration=self.configuration,
                 metadata=self.metadata,
@@ -79,6 +72,23 @@ def get_total(self, measure=None):
         totals_df = read_sql(totals_query[0], self.get_connection(), params=totals_query[1])
         return totals_df.iat[0, 0]
 
+    def get_row_count(self):
+        return self.__get_total__(None)
+
+    def get_total(self, measure=None):
+        """
+        For totals when streaming data we need to run a totals SQL query separate
+        from the main query and use the results of this
+        :return: the total value of the query
+        NOTE this is a direct copy from StreamingDataExtractor
+        """
+        if measure is None:
+            logger.warning(GET_TOTAL_WITHOUT_MEASURE)
+        measure = self.__get_measure__(measure)
+        if measure is None:
+            return self.get_row_count()
+        return self.__get_total__(measure)
+
     def run_query(self) -> str:
         """
         Runs the SQL query in Athena and returns the
@@ -190,4 +200,4 @@ def repl(match):
             raise KeyError(f"Missing SQL parameter: {name}")
         return athena_literal(params[name])
 
-    return ATHENA_PARAM_PATTERN.sub(repl, sql)
+    return ATHENA_PARAM_PATTERN.sub(repl, sql)

mario/data_extractor.py

@@ -12,6 +12,11 @@
 logger = logging.getLogger(__name__)
 
 
+GET_TOTAL_WITHOUT_MEASURE = 'Deprecation warning: Calling get_total() without any measure specified, or with None, ' \
+                            'is going to be deprecated in future. Instead you should either call get_row_count(), ' \
+                            'or get_total(measure) where the measure exists in the dataset.'
+
+
 class Configuration:
     """
     Configuration for a data extractor. This can include:
@@ -131,14 +136,35 @@ def __get_measure__(self, measure=None):
                 raise ValueError(f"Measure {measure} does not exist in dataset specification")
         return measure
 
-    def get_total(self, measure=None):
+    def get_row_count(self):
+        """
+        :return: the number of rows in the data
+        """
+        df = self.get_data_frame()
+        return len(df)
+
+    def __get_total__(self, measure):
         df = self.get_data_frame()
-        measure = self.__get_measure__(measure)
-        if measure is None:
-            return len(df)
         self._total = df[measure].sum()
         return self._total
 
+    def get_total(self, measure=None):
+        """
+        Returns the total from the dataset, using either the supplied measure, or
+        the default measure from the dataset specification if no measure is supplied.
+        If an invalid measure is supplied, a ValueError will be raised.
+        If no measure is supplied, and no measure exists in the dataset, the row_count()
+        is returned instead.
+        :param measure: the measure to use
+        :return: the sum of the data with the measure
+        """
+        if measure is None:
+            logger.warning(GET_TOTAL_WITHOUT_MEASURE)
+        measure = self.__get_measure__(measure)
+        if measure is None:
+            return self.get_row_count()
+        return self.__get_total__(measure)
+
     def validate_data(self, allow_nulls=True):
         from mario.validation import DataFrameValidator
         if self._data is None:
@@ -256,22 +282,34 @@ def __minimise_data__(self):
             table=table
         )
 
-    def get_total(self, measure=None):
-        from mario.hyper_utils import get_row_count, get_default_table_and_schema, get_total
+    def get_row_count(self):
+        from mario.hyper_utils import get_row_count as __get_row_count__
+        from mario.hyper_utils import get_default_table_and_schema
+        schema, table = get_default_table_and_schema(self.configuration.file_path)
+        return __get_row_count__ (
+            hyper_file_path=self.configuration.file_path,
+            schema=schema,
+            table=table
+        )
+
+    def __get_total__(self, measure):
+        from mario.hyper_utils import get_total as __get_total__
+        from mario.hyper_utils import get_default_table_and_schema
         schema, table = get_default_table_and_schema(self.configuration.file_path)
+        return  __get_total__(
+            hyper_file_path=self.configuration.file_path,
+            schema=schema,
+            table=table,
+            measure=measure
+        )
+
+    def get_total(self, measure=None):
         if measure is None:
-            return get_row_count(
-                hyper_file_path=self.configuration.file_path,
-                schema=schema,
-                table=table
-            )
-        else:
-            return get_total(
-                hyper_file_path=self.configuration.file_path,
-                schema=schema,
-                table=table,
-                measure=measure
-            )
+            logger.warning(GET_TOTAL_WITHOUT_MEASURE)
+        measure = self.__get_measure__(measure)
+        if measure is None:
+            return self.get_row_count()
+        return self.__get_total__(measure)
 
     def save_data_as_hyper(self, file_path: str, **kwargs):
         from mario.hyper_utils import save_hyper_as_hyper, add_row_numbers_to_hyper, get_default_table_and_schema
@@ -360,26 +398,35 @@ def save_data_as_hyper(self, file_path: str, **kwargs):
             **kwargs
         )
 
-    def get_total(self, measure=None):
-        """
-        For totals when streaming data we need to run a totals SQL query separate
-        from the main query and use the results of this
-        :return: the total value of the query
-        """
-        logger.info("Building totals query")
-        measure = self.__get_measure__(measure)
+    def __get_total__(self, measure):
         if self.configuration.query_builder is not None:
             from mario.query_builder import QueryBuilder
+            logger.info("Building totals query")
             query_builder: QueryBuilder = self.configuration.query_builder(
                 configuration=self.configuration,
                 metadata=self.metadata,
                 dataset_specification=self.dataset_specification)
             totals_query = query_builder.create_totals_query(measure=measure)
+            totals_df = pd.read_sql(totals_query[0], self.get_connection(), params=totals_query[1])
+            return totals_df.iat[0, 0]
         else:
             raise NotImplementedError
 
-        totals_df = pd.read_sql(totals_query[0], self.get_connection(), params=totals_query[1])
-        return totals_df.iat[0, 0]
+    def get_row_count(self):
+        return self.__get_total__(None)
+
+    def get_total(self, measure=None):
+        """
+        For totals when streaming data we need to run a totals SQL query separate
+        from the main query and use the results of this
+        :return: the total value of the query
+        """
+        if measure is None:
+            logger.warning(GET_TOTAL_WITHOUT_MEASURE)
+        measure = self.__get_measure__(measure)
+        if measure is None:
+            return self.get_row_count()
+        return self.__get_total__(measure)
 
     def stream_sql_to_hyper(self, file_path: str, **kwargs):
         """
@@ -609,6 +656,9 @@ def __load_from_sql__(self):
     def get_data_frame(self, minimise=True, include_row_numbers=False) -> DataFrame:
         raise NotImplementedError()
 
+    def get_row_count(self):
+        raise NotImplementedError()
+
     def get_total(self, measure=None):
         raise NotImplementedError()
 

mario/query_builder.py

@@ -101,23 +101,23 @@ def contains_subject(self) -> bool:
                     return True
         return False
 
+    def __get_measure_label__(self, measure:str):
+        if measure in ['FPE', 'FTE'] and not self.contains_subject():
+            return 'Count'
+        return measure
+
     def create_totals_query(self, measure=None) -> [str, List[any]]:
         """
         Constructs a 'total only' query by combining the selections and constraints.
         :return: an array containing the query prepared statement, and the parameters
-        TODO make this more generic
         """
         measures = []
-        measure = self.mapping.as_physical[measure]
         if measure is None:
-            for measure in self.dataset_specification.measures:
-                # Decide column name based on whether 'subject' fields are present
-                if measure in ['FPE', 'FTE'] and not self.contains_subject():
-                    measures.append(fn.Sum(self.table[measure], 'Count'))
-                else:
-                    measures.append(fn.Sum(self.table[measure], measure))
+            measures.append(fn.Count('*', 'Count'))
         else:
-            measures.append(fn.Sum(self.table[measure], measure))
+            measure_label = self.__get_measure_label__(measure)
+            measure = self.mapping.as_physical[measure]
+            measures.append(fn.Sum(self.table[measure], measure_label))
 
         q = Query().from_(self.table).select(*measures)
 
@@ -141,15 +141,12 @@ def create_query(self) -> [str, List[any]]:
         # remove measures from regular select
         measures = []
         for measure_field in self.dataset_specification.measures:
+            measure_label = self.__get_measure_label__(measure_field)
             measure = self.mapping.as_physical[measure_field]
             if measure in select_fields:
                 select_fields.remove(measure)
                 group_fields.remove(measure)
-                # Decide column name based on whether 'subject' fields are present
-                if measure in ['FPE', 'FTE'] and not self.contains_subject():
-                    measures.append(fn.Sum(self.table[measure], 'Count'))
-                else:
-                    measures.append(fn.Sum(self.table[measure], measure))
+                measures.append(fn.Sum(self.table[measure], measure_label))
         select_fields.extend(measures)
 
         q = Query(). \

test/test_athena_extractor.py

@@ -106,6 +106,26 @@ def test_athena_count():
     total = extractor.get_total(measure=dataset.measures[0])
     assert total == 28_733_910
 
+
+def test_athena_row_count():
+    if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION:
+        pytest.skip("Skipping Athena test as AWS not configured")
+
+    dataset, metadata, cfg = get_test_conf()
+    cfg.query_builder = SubsetQueryBuilder
+    cfg.schema = 'demo'
+    cfg.view = 'student_open_data'
+
+    extractor = AthenaDataExtractor(
+        configuration=cfg,
+        metadata=metadata,
+        dataset_specification=dataset
+    )
+
+    total = extractor.get_row_count()
+    assert total == 28_733_910
+
+
 def test_athena_count_with_space():
     if not AWS_PROFILE or not AWS_ATHENA_RESULTS_DIR or not AWS_REGION:
         pytest.skip("Skipping Athena test as AWS not configured")