Merge pull request #172 from PySport/feature/add-polars

Add polars support

Author: koenvo

kloppy/config.py

@@ -1,7 +1,7 @@
 import os
 from contextlib import contextmanager
 from copy import copy
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from kloppy.domain import EventFactory
 
@@ -27,6 +27,9 @@
         "event_factory": Optional[EventFactory],
         "adapters.http.basic_authentication": Optional[str],
         "adapters.s3.s3fs": Optional[Any],
+        "dataframe.engine": Optional[
+            Union[Literal["pandas"], Literal["polars"]]
+        ],
     },
 )
 
@@ -37,6 +40,7 @@
     "event_factory",
     "adapters.http.basic_authentication",
     "adapters.s3.s3fs",
+    "dataframe.engine",
 ]
 
 
@@ -50,6 +54,7 @@ class PartialConfig(Config, total=False):
     "event_factory": None,
     "adapters.http.basic_authentication": None,
     "adapters.s3.s3fs": None,
+    "dataframe.engine": "pandas",
 }
 
 config = copy(_default_config)

kloppy/domain/models/code.py

@@ -4,7 +4,10 @@
 from kloppy.domain.models.common import DatasetType
 
 from .common import Dataset, DataRecord
-from ...utils import docstring_inherit_attributes
+from kloppy.utils import (
+    docstring_inherit_attributes,
+    deprecated,
+)
 
 
 @dataclass
@@ -45,6 +48,9 @@ class CodeDataset(Dataset[Code]):
     def codes(self):
         return self.records
 
+    @deprecated(
+        "to_pandas will be removed in the future. Please use to_df instead."
+    )
     def to_pandas(
         self,
         record_converter: Callable[[Code], Dict] = None,
@@ -61,18 +67,11 @@ def to_pandas(
             )
 
         if not record_converter:
+            from ..services.transformers.attribute import (
+                DefaultCodeTransformer,
+            )
 
-            def record_converter(code: Code) -> Dict:
-                row = dict(
-                    code_id=code.code_id,
-                    period_id=code.period.id if code.period else None,
-                    timestamp=code.timestamp,
-                    end_timestamp=code.end_timestamp,
-                    code=code.code,
-                )
-                row.update(code.labels)
-
-                return row
+            record_converter = DefaultCodeTransformer()
 
         def generic_record_converter(code: Code):
             row = record_converter(code)

kloppy/domain/models/common.py

@@ -1,14 +1,33 @@
+import sys
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field, replace
 from enum import Enum, Flag
-from typing import Dict, List, Optional, Callable, Union, Any, TypeVar, Generic
+from typing import (
+    Dict,
+    List,
+    Optional,
+    Callable,
+    Union,
+    Any,
+    TypeVar,
+    Generic,
+    NewType,
+    overload,
+    Iterable,
+)
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
 
 from .pitch import PitchDimensions, Point, Dimension
 from .formation import FormationType
 from ...exceptions import (
     OrientationError,
-    OrphanedRecordError,
     InvalidFilterError,
+    KloppyParameterError,
 )
 
 
@@ -777,6 +796,8 @@ class Dataset(ABC, Generic[T]):
 
     """
 
+    Column = NewType("Column", Union[str, Callable[[T], Any]])
+
     records: List[T]
     metadata: Metadata
 
@@ -882,3 +903,65 @@ def get_record_by_id(self, record_id: Union[int, str]) -> Optional[T]:
         for record in self.records:
             if record.record_id == record_id:
                 return record
+
+    @overload
+    def to_records(
+        self,
+        *columns: "Column",
+        as_list: Literal[True] = True,
+        **named_columns: "Column",
+    ) -> List[Dict[str, Any]]:
+        ...
+
+    @overload
+    def to_records(
+        self,
+        *columns: "Column",
+        as_list: Literal[False] = False,
+        **named_columns: "Column",
+    ) -> Iterable[Dict[str, Any]]:
+        ...
+
+    def to_records(
+        self,
+        *columns: "Column",
+        as_list: bool = True,
+        **named_columns: "Column",
+    ) -> Union[List[Dict[str, Any]], Iterable[Dict[str, Any]]]:
+
+        from ..services.transformers.data_record import get_transformer_cls
+
+        transformer = get_transformer_cls(self.dataset_type)(
+            *columns, **named_columns
+        )
+        iterator = map(transformer, self.records)
+        if as_list:
+            return list(iterator)
+        else:
+            return iterator
+
+    def to_df(
+        self,
+        *columns: "Column",
+        engine: Optional[Union[Literal["polars"], Literal["pandas"]]] = None,
+        **named_columns: "Column",
+    ):
+        from kloppy.config import get_config
+
+        if not engine:
+            engine = get_config("dataframe.engine")
+
+        if engine == "pandas":
+            from pandas import DataFrame
+
+            return DataFrame.from_records(
+                self.to_records(*columns, **named_columns, as_list=False)
+            )
+        elif engine == "polars":
+            from polars import DataFrame
+
+            return DataFrame(
+                self.to_records(*columns, **named_columns, as_list=False)
+            )
+        else:
+            raise KloppyParameterError(f"Engine {engine} is not valid")

kloppy/domain/models/event.py

@@ -1,5 +1,3 @@
-import sys
-
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
@@ -11,21 +9,15 @@
     Any,
     Callable,
     Optional,
-    Iterable,
-    overload,
     TYPE_CHECKING,
 )
 
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
-
 from kloppy.domain.models.common import DatasetType
 from kloppy.utils import (
     camelcase_to_snakecase,
     removes_suffix,
     docstring_inherit_attributes,
+    deprecated,
 )
 
 from .common import DataRecord, Dataset, Player, Team
@@ -35,7 +27,6 @@
 from ...exceptions import OrphanedRecordError, InvalidFilterError
 
 if TYPE_CHECKING:
-    from ..services.transformers.event import Column
     from .tracking import Frame
 
 
@@ -819,6 +810,9 @@ def add_state(self, *builder_keys):
 
         return add_state(self, *builder_keys)
 
+    @deprecated(
+        "to_pandas will be removed in the future. Please use to_df instead."
+    )
     def to_pandas(
         self,
         record_converter: Callable[[Event], Dict] = None,
@@ -835,9 +829,11 @@ def to_pandas(
             )
 
         if not record_converter:
-            from ..services.transformers.attribute import DefaultTransformer
+            from ..services.transformers.attribute import (
+                DefaultEventTransformer,
+            )
 
-            record_converter = DefaultTransformer()
+            record_converter = DefaultEventTransformer()
 
         def generic_record_converter(event: Event):
             row = record_converter(event)
@@ -854,46 +850,6 @@ def generic_record_converter(event: Event):
             map(generic_record_converter, self.records)
         )
 
-    @overload
-    def to_records(
-        self,
-        *columns: "Column",
-        as_list: Literal[True] = True,
-        **named_columns: "Column",
-    ) -> List[Dict[str, Any]]:
-        ...
-
-    @overload
-    def to_records(
-        self,
-        *columns: "Column",
-        as_list: Literal[False] = False,
-        **named_columns: "Column",
-    ) -> Iterable[Dict[str, Any]]:
-        ...
-
-    def to_records(
-        self,
-        *columns: "Column",
-        as_list: bool = True,
-        **named_columns: "Column",
-    ) -> Union[List[Dict[str, Any]], Iterable[Dict[str, Any]]]:
-        from ..services.transformers.event import EventToRecordTransformer
-
-        transformer = EventToRecordTransformer(*columns, **named_columns)
-        iterator = map(transformer, self.events)
-        if as_list:
-            return list(iterator)
-        else:
-            return iterator
-
-    def to_df(self, *columns: "Column", **named_columns: "Column"):
-        from pandas import DataFrame
-
-        return DataFrame.from_records(
-            self.to_records(*columns, **named_columns, as_list=False)
-        )
-
 
 __all__ = [
     "EnumQualifier",

kloppy/domain/models/tracking.py

@@ -5,6 +5,9 @@
 
 from .common import Dataset, DataRecord, Player
 from .pitch import Point, Point3D
+from kloppy.utils import (
+    deprecated,
+)
 
 
 @dataclass
@@ -48,6 +51,9 @@ def frames(self):
     def frame_rate(self):
         return self.metadata.frame_rate
 
+    @deprecated(
+        "to_pandas will be removed in the future. Please use to_df instead."
+    )
     def to_pandas(
         self,
         record_converter: Callable[[Frame], Dict] = None,
@@ -64,59 +70,11 @@ def to_pandas(
             )
 
         if not record_converter:
+            from ..services.transformers.attribute import (
+                DefaultFrameTransformer,
+            )
 
-            def record_converter(frame: Frame) -> Dict:
-                row = dict(
-                    period_id=frame.period.id if frame.period else None,
-                    timestamp=frame.timestamp,
-                    ball_state=frame.ball_state.value
-                    if frame.ball_state
-                    else None,
-                    ball_owning_team_id=frame.ball_owning_team.team_id
-                    if frame.ball_owning_team
-                    else None,
-                    ball_x=frame.ball_coordinates.x
-                    if frame.ball_coordinates
-                    else None,
-                    ball_y=frame.ball_coordinates.y
-                    if frame.ball_coordinates
-                    else None,
-                    ball_z=getattr(frame.ball_coordinates, "z", None)
-                    if frame.ball_coordinates
-                    else None,
-                )
-                for player, player_data in frame.players_data.items():
-
-                    row.update(
-                        {
-                            f"{player.player_id}_x": player_data.coordinates.x
-                            if player_data.coordinates
-                            else None,
-                            f"{player.player_id}_y": player_data.coordinates.y
-                            if player_data.coordinates
-                            else None,
-                            f"{player.player_id}_d": player_data.distance,
-                            f"{player.player_id}_s": player_data.speed,
-                        }
-                    )
-
-                    if player_data.other_data:
-                        for name, value in player_data.other_data.items():
-                            row.update(
-                                {
-                                    f"{player.player_id}_{name}": value,
-                                }
-                            )
-
-                if frame.other_data:
-                    for name, value in frame.other_data.items():
-                        row.update(
-                            {
-                                name: value,
-                            }
-                        )
-
-                return row
+            record_converter = DefaultFrameTransformer()
 
         def generic_record_converter(frame: Frame):
             row = record_converter(frame)