diff --git a/gooddata-pandas/gooddata_pandas/result_convertor.py b/gooddata-pandas/gooddata_pandas/result_convertor.py index 6a2002318..c12881e0f 100644 --- a/gooddata-pandas/gooddata_pandas/result_convertor.py +++ b/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, diff --git a/gooddata-pandas/tests/dataframe/test_dataframe_for_exec_def.py b/gooddata-pandas/tests/dataframe/test_dataframe_for_exec_def.py index 4973292ac..f6ed2a686 100644 --- a/gooddata-pandas/tests/dataframe/test_dataframe_for_exec_def.py +++ b/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