diff --git a/frictionless/steps/table/__spec__/test_table_pivot.py b/frictionless/steps/table/__spec__/test_table_pivot.py index 6d720ff953..33a372b4ac 100644 --- a/frictionless/steps/table/__spec__/test_table_pivot.py +++ b/frictionless/steps/table/__spec__/test_table_pivot.py @@ -24,3 +24,33 @@ def test_step_table_pivot_issue_1220(): {"region": "east", "boy": 33, "girl": 29}, {"region": "west", "boy": 35, "girl": 23}, ] + + +def test_step_table_pivot_string_aggfun_issue_1764(): + source = TableResource(path="data/transform-pivot.csv") + pipeline = Pipeline.from_descriptor( + { + "steps": [ + {"type": "table-normalize"}, + { + "type": "table-pivot", + "f1": "region", + "f2": "gender", + "f3": "units", + "aggfun": "sum", + }, + ] + } + ) + target = source.transform(pipeline) + assert target.schema.to_descriptor() == { + "fields": [ + {"name": "region", "type": "string"}, + {"name": "boy", "type": "integer"}, + {"name": "girl", "type": "integer"}, + ] + } + assert target.read_rows() == [ + {"region": "east", "boy": 33, "girl": 29}, + {"region": "west", "boy": 35, "girl": 23}, + ] diff --git a/frictionless/steps/table/table_pivot.py b/frictionless/steps/table/table_pivot.py index b9e6bba28e..f11bc67255 100644 --- a/frictionless/steps/table/table_pivot.py +++ b/frictionless/steps/table/table_pivot.py @@ -1,9 +1,12 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +import statistics +from typing import TYPE_CHECKING, Any, Callable import attrs +from ... import errors +from ...exception import FrictionlessException from ...pipeline import Step from ...schema import Schema @@ -11,6 +14,21 @@ from ...resource import Resource +# Mapping of petl aggregation function names to callables. This allows a +# string `aggfun` (e.g. coming from a JSON/YAML descriptor) to mirror the +# Python API, where a callable is passed directly. +AGGREGATION_FUNCTIONS: dict[str, Callable[[Any], Any]] = { + "sum": sum, + "max": max, + "min": min, + "len": len, + "count": len, + "mean": statistics.mean, + "first": lambda values: list(values)[0], + "last": lambda values: list(values)[-1], +} + + @attrs.define(kw_only=True, repr=False) class table_pivot(Step): """Pivot table. @@ -40,14 +58,22 @@ class table_pivot(Step): aggfun: Any """ Function to process and create data in the output pivot table. - The function can be "sum", "max", "min", "len" etc. + It can be a callable or, when defined via a descriptor, one of the + supported aggregation function names: "sum", "max", "min", "len", + "count", "mean", "first", "last". """ # Transform def transform_resource(self, resource: Resource): table = resource.to_petl() # type: ignore - resource.data = table.pivot(self.f1, self.f2, self.f3, self.aggfun) # type: ignore + aggfun = self.aggfun + if isinstance(aggfun, str): + if aggfun not in AGGREGATION_FUNCTIONS: + note = f'aggregation function "{aggfun}" is not supported' + raise FrictionlessException(errors.StepError(note=note)) + aggfun = AGGREGATION_FUNCTIONS[aggfun] + resource.data = table.pivot(self.f1, self.f2, self.f3, aggfun) # type: ignore resource.schema = Schema() resource.infer()