Merge pull request #31 from rjtavares/pandas_additional_columns

added option to parse additional columns when calling to_pandas

README.md

@@ -247,6 +247,21 @@ new_dataset = Transformer.transform_dataset(
 )
 ```
 
+### Adding additional columns to DataFrame
+
+You can add more information to the DataFrame returned by `to_pandas` by passing a Dictionary in the format `{'column_name': function or scalar}`. Function should take a frame (for tracking data) or an event (for event data) as argument and return a valid input for a DataFrame.
+
+```python
+from kloppy import datasets, to_pandas
+dataset = datasets.load("statsbomb")
+
+# Statsbomb events may include a list of related events.
+# This function returns the first event from that list, if it exists.
+get_first_related_event = lambda x: x.raw_event['related_events'][0] if 'related_events' in x.raw_event else None
+to_pandas(dataset, additional_columns={'related_event': get_first_related_event, 'match': 'test1'})
+```
+
+
 # Contributing to kloppy
 All contributions, bug reports, bug fixes, documentation improvements, enhancements, and ideas are welcome.
 

kloppy/helpers.py

@@ -1,4 +1,4 @@
-from typing import Callable, TypeVar, Dict
+from typing import Callable, TypeVar, Dict, Union
 
 from . import (
     TRACABSerializer,
@@ -193,7 +193,7 @@ def _event_to_pandas_row_converter(event: Event) -> Dict:
 
 
 def to_pandas(
-    dataset: Dataset, _record_converter: Callable = None
+    dataset: Dataset, _record_converter: Callable = None, additional_columns: Dict = None
 ) -> "DataFrame":
     try:
         import pandas as pd
@@ -211,7 +211,19 @@ def to_pandas(
         else:
             raise Exception("Unknown dataset type")
 
-    return pd.DataFrame.from_records(map(_record_converter, dataset.records))
+    def generic_record_converter(record: Union[Frame, Event]):
+        row = _record_converter(record)
+        if additional_columns:
+            for k, v in additional_columns.items():
+                if callable(v):
+                    value = v(record)
+                else:
+                    value = v
+                row.update({k: value})
+
+        return row
+
+    return pd.DataFrame.from_records(map(generic_record_converter, dataset.records))
 
 
 __all__ = [