Merge pull request #1128 from Martozar/c.mze-cq-1474

feat: add column total header indexes to DataFrameMetadata

Author: Martozar

gooddata-pandas/gooddata_pandas/result_convertor.py

@@ -437,11 +437,13 @@ class DataFrameMetadata:
                                     | AVG     |    150
                           SUM       |         |    450
 
+      column_totals_indexes: Similar to row_totals_indexes but for column headers.
       execution_response: An instance of BareExecutionResponse representing the
                           execution response.
     """
 
     row_totals_indexes: list[list[int]]
+    column_totals_indexes: list[list[int]]
     execution_response: BareExecutionResponse
     primary_labels_from_index: dict[int, dict[str, str]]
     primary_labels_from_columns: dict[int, dict[str, str]]
@@ -460,17 +462,25 @@ def from_data(
             A tuple containing data headers. execution_response (BareExecutionResponse): An ExecutionResponse object.
 
         Returns: DataFrameMetadata: An initialized DataFrameMetadata object."""
-        row_totals_indexes = [
-            [idx for idx, hdr in enumerate(dim) if hdr is not None and hdr.get("totalHeader") is not None]
-            for dim in headers[0]
-        ]
+        row_totals_indexes = cls._get_totals_indexes(headers[0])
+        column_totals_indexes = cls._get_totals_indexes(headers[1])
         return cls(
             row_totals_indexes=row_totals_indexes,
+            column_totals_indexes=column_totals_indexes,
             execution_response=execution_response,
             primary_labels_from_index=primary_labels_from_index,
             primary_labels_from_columns=primary_labels_from_columns,
         )
 
+    @staticmethod
+    def _get_totals_indexes(headers: Optional[Any]) -> list[list[int]]:
+        if headers is None:
+            return []
+        return [
+            [idx for idx, hdr in enumerate(dim) if hdr is not None and hdr.get("totalHeader") is not None]
+            for dim in headers
+        ]
+
 
 def _read_complete_execution_result(
     execution_response: BareExecutionResponse,

gooddata-pandas/tests/dataframe/test_dataframe_for_exec_def.py

@@ -28,6 +28,7 @@ def _run_and_validate_results(
     exec_def: ExecutionDefinition,
     expected: tuple[int, int],
     expected_row_totals: Optional[list[list[int]]] = None,
+    expected_column_totals: Optional[list[list[int]]] = None,
     page_size: int = 100,
     optimized: bool = False,
 ) -> str:
@@ -42,6 +43,8 @@ def _run_and_validate_results(
 
     if expected_row_totals is not None:
         assert result_metadata_from_result_id.row_totals_indexes == expected_row_totals
+    if expected_column_totals is not None:
+        assert result_metadata_from_result_id.column_totals_indexes == expected_column_totals
 
     assert result_from_result_id.values.shape == expected
 
@@ -349,7 +352,9 @@ def test_dataframe_for_exec_def_totals4(gdf: DataFrameFactory, optimized: bool):
             ),
         ],
     )
-    _run_and_validate_results(gdf=gdf, exec_def=exec_def, expected=(96, 19), optimized=optimized)
+    _run_and_validate_results(
+        gdf=gdf, exec_def=exec_def, expected=(96, 19), optimized=optimized, expected_column_totals=[[17, 18], [17, 18]]
+    )
 
 
 # TODO - not implemented yet