Skip to content

Commit 823cd2b

Browse files
committed
feat: add column total header indexes to DataFrameMetadata
JIRA: CQ-1474 risk: low
1 parent 4f27fa5 commit 823cd2b

File tree

2 files changed

+19
-5
lines changed

2 files changed

+19
-5
lines changed

gooddata-pandas/gooddata_pandas/result_convertor.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -437,6 +437,7 @@ class DataFrameMetadata:
437437
| AVG | 150
438438
SUM | | 450
439439
440+
column_totals_indexes: Similar to row_totals_indexes but for column headers.
440441
execution_response: An instance of BareExecutionResponse representing the
441442
execution response.
442443
"""
@@ -460,17 +461,25 @@ def from_data(
460461
A tuple containing data headers. execution_response (BareExecutionResponse): An ExecutionResponse object.
461462
462463
Returns: DataFrameMetadata: An initialized DataFrameMetadata object."""
463-
row_totals_indexes = [
464-
[idx for idx, hdr in enumerate(dim) if hdr is not None and hdr.get("totalHeader") is not None]
465-
for dim in headers[0]
466-
]
464+
row_totals_indexes = cls._get_totals_indexes(headers[0])
465+
column_totals_indexes = cls._get_totals_indexes(headers[1])
467466
return cls(
468467
row_totals_indexes=row_totals_indexes,
468+
column_totals_indexes=column_totals_indexes,
469469
execution_response=execution_response,
470470
primary_labels_from_index=primary_labels_from_index,
471471
primary_labels_from_columns=primary_labels_from_columns,
472472
)
473473

474+
@staticmethod
475+
def _get_totals_indexes(headers: Optional[Any]) -> list[list[int]]:
476+
if headers is None:
477+
return []
478+
return [
479+
[idx for idx, hdr in enumerate(dim) if hdr is not None and hdr.get("totalHeader") is not None]
480+
for dim in headers
481+
]
482+
474483

475484
def _read_complete_execution_result(
476485
execution_response: BareExecutionResponse,

gooddata-pandas/tests/dataframe/test_dataframe_for_exec_def.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ def _run_and_validate_results(
2828
exec_def: ExecutionDefinition,
2929
expected: tuple[int, int],
3030
expected_row_totals: Optional[list[list[int]]] = None,
31+
expected_column_totals: Optional[list[list[int]]] = None,
3132
page_size: int = 100,
3233
optimized: bool = False,
3334
) -> str:
@@ -42,6 +43,8 @@ def _run_and_validate_results(
4243

4344
if expected_row_totals is not None:
4445
assert result_metadata_from_result_id.row_totals_indexes == expected_row_totals
46+
if expected_column_totals is not None:
47+
assert result_metadata_from_result_id.column_totals_indexes == expected_column_totals
4548

4649
assert result_from_result_id.values.shape == expected
4750

@@ -349,7 +352,9 @@ def test_dataframe_for_exec_def_totals4(gdf: DataFrameFactory, optimized: bool):
349352
),
350353
],
351354
)
352-
_run_and_validate_results(gdf=gdf, exec_def=exec_def, expected=(96, 19), optimized=optimized)
355+
_run_and_validate_results(
356+
gdf=gdf, exec_def=exec_def, expected=(96, 19), optimized=optimized, expected_column_totals=[[17, 18], [17, 18]]
357+
)
353358

354359

355360
# TODO - not implemented yet

0 commit comments

Comments
 (0)